fix: Fixing codebase agent

src/codegen/extensions/langchain/agent.py

@@ -3,7 +3,7 @@
 from langchain import hub
 from langchain.agents import AgentExecutor
 from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
-from langchain_core.chat_history import ChatMessageHistory
+from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_openai import ChatOpenAI
 
@@ -20,6 +20,7 @@
     RevealSymbolTool,
     SearchTool,
     SemanticEditTool,
+    SemanticSearchTool,
     ViewFileTool,
 )
 
@@ -59,6 +60,7 @@ def create_codebase_agent(
         MoveSymbolTool(codebase),
         RevealSymbolTool(codebase),
         SemanticEditTool(codebase),
+        SemanticSearchTool(codebase),
         CommitTool(codebase),
     ]
 

src/codegen/extensions/langchain/tools.py

@@ -33,9 +33,9 @@
 class ViewFileTool(BaseTool):
     """Tool for viewing file contents and metadata."""
 
     name: ClassVar[str] = "view_file"
     description: ClassVar[str] = "View the contents and metadata of a file in the codebase"
     args_schema: ClassVar[type[BaseModel]] = ViewFileInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -56,9 +56,9 @@
 class ListDirectoryTool(BaseTool):
     """Tool for listing directory contents."""
 
     name: ClassVar[str] = "list_directory"
     description: ClassVar[str] = "List contents of a directory in the codebase"
     args_schema: ClassVar[type[BaseModel]] = ListDirectoryInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -79,9 +79,9 @@
 class SearchTool(BaseTool):
     """Tool for searching the codebase."""
 
     name: ClassVar[str] = "search"
     description: ClassVar[str] = "Search the codebase using text search"
     args_schema: ClassVar[type[BaseModel]] = SearchInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -102,7 +102,7 @@
 class EditFileTool(BaseTool):
     """Tool for editing files."""
 
     name: ClassVar[str] = "edit_file"
     description: ClassVar[str] = "Edit a file by replacing its entire content"
     args_schema: ClassVar[type[BaseModel]] = EditFileInput
     codebase: Codebase = Field(exclude=True)
@@ -312,20 +312,18 @@
         return json.dumps(result, indent=2)
 
 
+class SemanticSearchInput(BaseModel):
+    query: str = Field(..., description="The natural language search query")
+    k: int = Field(default=5, description="Number of results to return")
+    preview_length: int = Field(default=200, description="Length of content preview in characters")
+
+
 class SemanticSearchTool(BaseTool):
     """Tool for semantic code search."""
 
     name: ClassVar[str] = "semantic_search"
     description: ClassVar[str] = "Search the codebase using natural language queries and semantic similarity"
-    args_schema: ClassVar[type[BaseModel]] = type(
-        "SemanticSearchInput",
-        (BaseModel,),
-        {
-            "query": (str, Field(..., description="The natural language search query")),
-            "k": (int, Field(default=5, description="Number of results to return")),
-            "preview_length": (int, Field(default=200, description="Length of content preview in characters")),
-        },
-    )
+    args_schema: ClassVar[type[BaseModel]] = SemanticSearchInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None: