update agents

app/modules/conversations/conversation/conversation_service.py

@@ -168,10 +168,15 @@ async def classifier_node(self, state: State) -> Command:
                 update={"response": "Error in classification format"}, goto=END
             )
         if confidence < 0.5 or agent_id not in self.agents:
+            logger.info(
+                f"Streaming AI response for conversation {state['conversation_id']} for user {state['user_id']} using agent {agent_id}"
+            )
             return Command(
                 update={"agent_id": state["agent_id"]}, goto=state["agent_id"]
             )
-
+        logger.info(
+                f"Streaming AI response for conversation {state['conversation_id']} for user {state['user_id']} using agent {agent_id}"
+            )
         return Command(update={"agent_id": agent_id}, goto=agent_id)
 
     async def agent_node(self, state: State, writer: StreamWriter):

app/modules/intelligence/agents/agents/blast_radius_agent.py

@@ -37,6 +37,11 @@ async def create_agents(self):
             role="Blast Radius Agent",
             goal="Explain the blast radius of the changes made in the code.",
             backstory="You are an expert in understanding the impact of code changes on the codebase.",
+            tools=[
+                get_change_detection_tool(self.user_id),
+                self.get_nodes_from_tags,
+                self.ask_knowledge_graph_queries,
+            ],
             allow_delegation=False,
             verbose=True,
             llm=self.llm,
@@ -106,11 +111,6 @@ async def create_tasks(
             {self.BlastRadiusAgentResponse.model_json_schema()}""",
             expected_output=f"Comprehensive impact analysis of the code changes on the codebase and answers to the users query about them. Ensure that your output ALWAYS follows the structure outlined in the following pydantic model : {self.BlastRadiusAgentResponse.model_json_schema()}",
             agent=blast_radius_agent,
-            tools=[
-                get_change_detection_tool(self.user_id),
-                self.get_nodes_from_tags,
-                self.ask_knowledge_graph_queries,
-            ],
             output_pydantic=self.BlastRadiusAgentResponse,
             async_execution=True,
         )

app/modules/intelligence/agents/agents/unit_test_agent.py

@@ -19,20 +19,26 @@
 
 class UnitTestAgent:
     def __init__(self, sql_db, llm, user_id):
+        self.openai_api_key = os.getenv("OPENAI_API_KEY")
+        self.max_iterations = os.getenv("MAX_ITER", 15)
         self.sql_db = sql_db
+        self.llm = llm
         self.user_id = user_id
+        # Initialize tools with both sql_db and user_id
         self.get_code_from_node_id = get_code_from_node_id_tool(sql_db, user_id)
         self.get_code_from_probable_node_name = get_code_from_probable_node_name_tool(
             sql_db, user_id
         )
-        self.llm = llm
-        self.max_iterations = os.getenv("MAX_ITER", 15)
 
     async def create_agents(self):
         unit_test_agent = Agent(
             role="Test Plan and Unit Test Expert",
             goal="Create test plans and write unit tests based on user requirements",
             backstory="You are a seasoned AI test engineer specializing in creating robust test plans and unit tests. You aim to assist users effectively in generating and refining test plans and unit tests, ensuring they are comprehensive and tailored to the user's project requirements.",
+            tools=[
+                self.get_code_from_probable_node_name,
+                self.get_code_from_node_id
+            ],
             allow_delegation=False,
             verbose=True,
             llm=self.llm,
@@ -117,7 +123,6 @@ async def create_tasks(
             expected_output="Outline the test plan and write unit tests for each node based on the test plan.",
             agent=unit_test_agent,
             output_pydantic=self.TestAgentResponse,
-            tools=[self.get_code_from_probable_node_name, self.get_code_from_node_id],
             async_execution=True,
         )
 

app/modules/intelligence/tools/kg_based_tools/get_code_from_multiple_node_ids_tool.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Dict, List
 
-from langchain.tools import StructuredTool
+from langchain_core.tools import StructuredTool, Tool    
 from neo4j import GraphDatabase
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
@@ -156,7 +156,7 @@ def __del__(self):
 
 def get_code_from_multiple_node_ids_tool(
     sql_db: Session, user_id: str
-) -> StructuredTool:
+) -> Tool:
     tool_instance = GetCodeFromMultipleNodeIdsTool(sql_db, user_id)
     return StructuredTool.from_function(
         coroutine=tool_instance.arun,

app/modules/intelligence/tools/kg_based_tools/get_code_from_node_id_tool.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Dict
 
-from langchain.tools import StructuredTool
+from langchain_core.tools import StructuredTool, Tool
 from neo4j import GraphDatabase
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
@@ -135,7 +135,7 @@ def __del__(self):
             self.neo4j_driver.close()
 
 
-def get_code_from_node_id_tool(sql_db: Session, user_id: str) -> StructuredTool:
+def get_code_from_node_id_tool(sql_db: Session, user_id: str) -> Tool:
     tool_instance = GetCodeFromNodeIdTool(sql_db, user_id)
     return StructuredTool.from_function(
         coroutine=tool_instance.arun,

app/modules/intelligence/tools/kg_based_tools/get_code_from_probable_node_name_tool.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Dict, List
 
-from langchain.tools import StructuredTool
+from langchain_core.tools import StructuredTool, Tool
 from neo4j import GraphDatabase
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
@@ -207,7 +207,7 @@ def __del__(self):
 
 def get_code_from_probable_node_name_tool(
     sql_db: Session, user_id: str
-) -> StructuredTool:
+) -> Tool:
     tool_instance = GetCodeFromProbableNodeNameTool(sql_db, user_id)
     return StructuredTool.from_function(
         coroutine=tool_instance.arun,

requirements.txt

@@ -52,5 +52,5 @@ sentry-sdk[fastapi]==2.19.0
 posthog==3.7.4
 newrelic==9.0.0
 tiktoken==0.7.0
-agentops==0.3.18
+agentops==0.3.23
 pydantic==2.10.3