feat: introduce plan and route with supervisor #140

muralov · 2024-08-20T08:15:15Z

Description
Introduces first implementation of Plan and Route features of Supervisor Agent.

Changes proposed in this pull request:

Implements supervisor agent
Introduces plan and route feature for supervisor agent
Introduces common node to answer general questions
Introduces finalizer node to generate final response
Introduces generate LangGraph diagram script
Implements unit tests

Related issue(s)
#112

* change it to /conversations/{conversation_id}/messages * Refactor some model code

* Build plan & route logic with langgraph for k8s and kyma agents

* Implement a langgraph node which summarizes the agent conversation as a final response * Upgrade langgraph

Main Kyma graph will have all the workflow of the companion. It'll have init node too.

* Write unit tests * Improve supervisor prompt to consider subtask statuses * Implement script to generate langgraph diagramm

* Create tests for all agents * Create tests for util functions * Create tests for kyma graph * Update the existing tests

* As we have only kyma and k8s questions deny irrelevant questions * Add coversation/id/messages route rest tests

* Improve plan & route with k8s and kyma terminology to route properly * Add common node to answer general questions * Improve error handling

mfaizanse · 2024-08-29T08:31:23Z

scripts/python/generate_langgraph_diagram.py

+
+from agents.graph import KymaGraph
+
+load_dotenv()


Can you write comments on how to run this script and what env variables are needed to run this script?

src/agents/common/agent.py

src/agents/common/state.py

mfaizanse · 2024-08-29T08:41:04Z

src/agents/common/state.py

+
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    next: str | None
+    subtasks: list[SubTask] | None = []


nitpick: sub_tasks?

let's leave as it is as the word subtask really exists :)

mfaizanse · 2024-08-29T08:41:32Z

src/agents/common/state.py

+class Plan(BaseModel):
+    """Plan to follow in future"""
+
+    subtasks: list[SubTask] = Field(


nitpick: sub_tasks?

src/agents/common/utils.py

mfaizanse · 2024-08-29T08:44:34Z

supervisor_agent.py

@@ -0,0 +1,233 @@
+import functools


Is this file supposed to be merged to main?
If not, please delete it.
If yes, please move to /examples or /hack directory.

mfaizanse · 2024-08-29T08:48:19Z

src/agents/supervisor/agent.py

+
+    model: Model
+    _name: str = SUPERVISOR
+    parser = PydanticOutputParser(pydantic_object=Plan)  # type: ignore


Is this parser used anywhere?

yes, below, you can search for it with state.subtasks = self.parser.parse(plan.content).subtasks

src/agents/graph.py

mfaizanse · 2024-08-29T11:15:15Z

src/agents/graph.py

+                ],
+            }
+
+    def common_node(self, state: AgentState) -> dict[str, Any]:


Can you explain what is the purpose of this node?

it answer the general questions. We'll move this to planner if possible in the next PR.

src/agents/graph.py

scripts/python/generate_langgraph_diagram.py

src/agents/common/utils.py

Teneroy · 2024-08-29T09:44:42Z

src/agents/graph.py

+PLANNER = "Planner"
+
+
+def should_continue(state: MessagesState) -> Literal["action", "__end__"]:


Why do we have should_exit in the utils, and should_continue here? Shouldn't it be moved to utils as well?

I'll delete it

Teneroy · 2024-08-29T10:10:25Z

src/agents/graph.py

+class CustomJSONEncoder(json.JSONEncoder):
+    """Custom JSON encoder."""
+
+    def default(self, obj):  # noqa D102


Intermediate steps have to be streamed correctly according to an agreement in our API documentation. Currently, we receive everything in our output

yes. There are mainly two reasons I haven't done it yet:

first, didn't have enough time

second, we need to add status code (error or success) for the chunks as anytime error may occur during streaming. The status code in the begining is almost always 200 due asynchronous behavior of FastAPI routes.

For that reason, I'd suggest to have follow-up ticket for this? We can think how to add the code and implement it in the follow-up ticket.
I'll also document this in the issue.

src/agents/graph.py

Teneroy · 2024-08-29T10:55:43Z

src/agents/graph.py

+    def common_node(self, state: AgentState) -> dict[str, Any]:
+        """Breaks down the given user query into sub-tasks."""
+
+        prompt = ChatPromptTemplate.from_messages(


Same thing as the commend above about the common node. Consider using history during the planning. Consider merging those two nodes

We need think whether we should consider the conversation or not for the planning. I am currently not sure.
I will document the points in the issue and improve them further in the next sprints.
Let's keep it simple for the start?

I thought we discussed this already, I've avoided this re-planning by separating out the.

if common questions, don't plan and directly send to Finalizer!

src/agents/supervisor/agent.py

src/services/conversation.py

Teneroy · 2024-08-29T11:33:32Z

src/utils/langfuse.py

+from langfuse.callback import CallbackHandler
+
+handler = CallbackHandler(
+    secret_key=os.getenv("LANGFUSE_SECRET_KEY"),


Let's move to some constant variables

src/agents/common/utils.py

Teneroy · 2024-08-29T11:49:29Z

src/agents/graph.py

+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                MessagesPlaceholder(variable_name="messages"),


Supervisor agent should have an ability to fail properly in case other agents can't provide relevant information. Therefore it should probably analyze provided information and make a decision whether this information is sufficient and it can generate a comprehensive response or it should fail. I suspect, additional prompting should be done in the finalizer. Here is an example where an agent fails in doing that:
request_response1.txt

In this case there 4 subtasks are created how to analyse the failing pod. We should discuss and optimize the planner accordingly in the next sprints. For that we need to first discuss whether we need to have detailed plan or keep the original sentences.

With this ticket, we know how we can do the planning and routing now. I'd suggest to optimize it step-by-step.

I'll document this too.

Teneroy · 2024-08-29T12:16:03Z

src/agents/supervisor/agent.py

-
-    result = agent.invoke(state)
-    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+logger = get_logger(__name__)


 class SupervisorAgent:


Consider renaming to task executor. Then we will merge planner, task executor, replan, and finalizer into the subgraph called SupervisorAgent in the future

I would leave this as SupervisorAgent as it eventually will have more functionality. Planning is just temporarily separated out for the start. It will be part of SupervisorAgent too soon.
It is currently not executing any task, it is looking at the conversation and subtasks and managing the communication to all other members of the supervisor team.

mfaizanse · 2024-08-30T13:35:56Z

src/agents/common/utils.py

+    }
+
+
+DEFAULT_NUMBER_OF_MESSAGES = 10


lets move this to the common file?

mfaizanse · 2024-08-30T13:43:13Z

src/services/conversation.py

+
+logger = get_logger(__name__)
+
+REDIS_URL = f"{os.getenv('REDIS_URL')}/0"


Ideally, this should be part of a separate module Config.

yes, I already have this in mind for the next PR. I will create a TODO for this too.

mfaizanse · 2024-08-30T13:44:59Z

src/services/conversation.py

+class SingletonMeta(type):
+    """Singleton metaclass."""
+
+    _instances: dict[type, object] = {}


Why it needs to be a dict? Can it be a single _instance: object?

this meta class have class field _instances which can be used by the other classes too. It means it stores single instances for multiple classes.

muralov added 2 commits August 15, 2024 15:13

Change chat route

aa87f57

* change it to /conversations/{conversation_id}/messages * Refactor some model code

Implement Plan & Route Logic

788d3e7

* Build plan & route logic with langgraph for k8s and kyma agents

kyma-bot added do-not-merge/work-in-progress cla: yes size/XXL labels Aug 20, 2024

Improve plan and route with Pydantic

296474c

muralov force-pushed the build-plan-and-route branch from 767247e to 296474c Compare August 20, 2024 10:05

muralov changed the title ~~Build plan and route~~ feat: introduce plan and route with supervisor Aug 20, 2024

muralov added 8 commits August 21, 2024 10:34

Implement final response node

c8f14a5

* Implement a langgraph node which summarizes the agent conversation as a final response * Upgrade langgraph

Make plan a separate node

af6e240

Create a main Kyma Graph

af73efc

Main Kyma graph will have all the workflow of the companion. It'll have init node too.

Improve Supervisor Prompt and Write Tests

e139f2c

* Write unit tests * Improve supervisor prompt to consider subtask statuses * Implement script to generate langgraph diagramm

Crete Unit tests

cb4ac8d

* Create tests for all agents * Create tests for util functions * Create tests for kyma graph * Update the existing tests

Deny non-kyma and non-k8s questions

113092c

* As we have only kyma and k8s questions deny irrelevant questions * Add coversation/id/messages route rest tests

Improve planning and routing prompts

dff9f0f

* Improve plan & route with k8s and kyma terminology to route properly * Add common node to answer general questions * Improve error handling

Add more to state and use FastAPI Depends

64ace79

muralov marked this pull request as ready for review August 29, 2024 08:02

muralov requested a review from a team as a code owner August 29, 2024 08:02

muralov requested a review from the1bit August 29, 2024 08:02

kyma-bot removed the do-not-merge/work-in-progress label Aug 29, 2024

muralov force-pushed the build-plan-and-route branch from 2aab769 to 71feb75 Compare August 29, 2024 08:11

muralov requested review from mfaizanse and removed request for the1bit August 29, 2024 08:14

Fix linting errors

ec98aee

muralov force-pushed the build-plan-and-route branch from 71feb75 to ec98aee Compare August 29, 2024 08:25