Merge pull request #77 from NASA-IMPACT/feature/combined-code-agent

Add CombinedCodeSearchTool for unified code search

Author: NISH1001

akd/tools/code_search.py

@@ -12,6 +12,7 @@
 from pydantic.networks import HttpUrl
 from scipy.spatial.distance import cdist
 from tenacity import retry, stop_after_attempt
+from sentence_transformers import CrossEncoder
 
 from akd.errors import SchemaValidationError
 from akd.structures import SearchResultItem
@@ -31,7 +32,10 @@ class CodeSearchToolInputSchema(SearchToolInputSchema):
     Input schema for the code search tool.
     """
 
-    pass
+    @computed_field
+    def top_k(self) -> int:
+        """Returns the number of top results to return."""
+        return self.max_results
 
 
 class CodeSearchToolOutputSchema(SearchToolOutputSchema):
@@ -120,11 +124,11 @@ def __get_sort_key(result):
 
             # Then check if 'extra' field exists and contains the sort_by key
             if (
-                "extra" in result
-                and isinstance(result["extra"], dict)
-                and sort_by in result["extra"]
+                result.extra
+                and isinstance(result.extra, dict)
+                and sort_by in result.extra
             ):
-                return result["extra"][sort_by]
+                return result.extra[sort_by]
 
             # If key not found anywhere, return a default value that will sort last
             # Using float('inf') for numerical sorting or empty string for string sorting
@@ -139,15 +143,85 @@ def __get_sort_key(result):
             return results
 
 
-class LocalRepoCodeSearchToolInputSchema(CodeSearchToolInputSchema):
+class CombinedCodeSearchToolConfig(CodeSearchToolConfig):
     """
-    Input schema for the local repository code search tool.
+    Configuration for the combined code search tool.
     """
 
-    @computed_field
-    def top_k(self) -> int:
-        """Returns the number of top results to return."""
-        return self.max_results
+    reranker_model_name: str = Field(
+        "cross-encoder/ms-marco-MiniLM-L6-v2",
+        description="The model to use for reranking the combined results.",
+    )
+
+
+class CombinedCodeSearchTool(CodeSearchTool):
+    """
+    Tool for performing combined code search using multiple sub-tools and CrossEncoder reranking.
+    """
+
+    input_schema = CodeSearchToolInputSchema
+    output_schema = CodeSearchToolOutputSchema
+    config_schema = CombinedCodeSearchToolConfig
+
+    def __init__(
+        self,
+        config: CombinedCodeSearchToolConfig | None = None,
+        tools: Optional[list[CodeSearchTool]] = None,
+        debug: bool = False,
+    ):
+        super().__init__(config, debug)
+        self.tools = tools or [
+            LocalRepoCodeSearchTool(debug=debug),
+            GitHubCodeSearchTool(debug=debug),
+            SDECodeSearchTool(debug=debug),
+        ]
+        self.reranker_model = CrossEncoder(config.reranker_model_name)
+
+    async def _arun(
+        self,
+        params: CodeSearchToolInputSchema,
+        **kwargs,
+    ) -> CodeSearchToolOutputSchema:
+        """Run the combined code search tool and aggregate results from all tools."""
+        all_results: list[SearchResultItem] = []
+
+        for tool in self.tools:
+            try:
+                if self.debug:
+                    logger.debug(f"Running tool: {tool.__class__.__name__}")
+                result = await tool._arun(params)
+                for res in result.results:
+                    res.extra["tool"] = tool.__class__.__name__
+                all_results.extend(result.results)
+            except Exception as e:
+                logger.error(f"Error running tool {tool.__class__.__name__}: {e}")
+
+        reranked = self._rerank_results(all_results, params.queries[0])[: params.top_k]
+        return self.output_schema(results=reranked, category="technology")
+
+    def _rerank_results(
+        self, results: list[SearchResultItem], query: str
+    ) -> list[SearchResultItem]:
+        """
+        Rerank results using a CrossEncoder model based on the query and result content.
+        """
+
+        try:
+            # Generate (query, content) pairs
+            pairs = [(query, result.content) for result in results]
+
+            # Get similarity scores from CrossEncoder
+            scores = self.reranker_model.predict(pairs)
+
+            # Attach scores
+            for score, result in zip(scores, results):
+                result.extra["score"] = score
+
+            # Sort results by score
+            return self._sort_results(results, sort_by="score")
+        except Exception as e:
+            logger.error(f"Reranking failed: {e}")
+            return results
 
 
 class LocalRepoCodeSearchToolConfig(CodeSearchToolConfig):
@@ -173,7 +247,7 @@ class LocalRepoCodeSearchTool(CodeSearchTool):
     It automatically downloads the necessary data file if it's not found locally.
     """
 
-    input_schema = LocalRepoCodeSearchToolInputSchema
+    input_schema = CodeSearchToolInputSchema
     output_schema = CodeSearchToolOutputSchema
     config_schema = LocalRepoCodeSearchToolConfig
 

examples/code_search_test.py

@@ -10,10 +10,11 @@
     CodeSearchToolInputSchema,
     LocalRepoCodeSearchTool,
     LocalRepoCodeSearchToolConfig,
-    LocalRepoCodeSearchToolInputSchema,
     GitHubCodeSearchTool,
     SDECodeSearchTool,
     SDECodeSearchToolConfig,
+    CombinedCodeSearchTool,
+    CombinedCodeSearchToolConfig,
 )
 
 
@@ -25,9 +26,7 @@ async def local_repo_search_test():
     cfg = LocalRepoCodeSearchToolConfig()
     tool = LocalRepoCodeSearchTool(config=cfg)
 
-    search_input = LocalRepoCodeSearchToolInputSchema(
-        queries=["landslide nepal"], max_results=5
-    )
+    search_input = CodeSearchToolInputSchema(queries=["landslide nepal"], max_results=5)
 
     print("Running the search...")
     output = await tool._arun(search_input)
@@ -82,10 +81,36 @@ async def sde_search_test():
         print("-" * 100)
 
 
+# Combined Code Search Tool
+async def combined_code_search_test():
+    """An async function to run the tool."""
+
+    print("Initializing the tool...")
+    cfg = CombinedCodeSearchToolConfig()
+    tool = CombinedCodeSearchTool(config=cfg)
+
+    search_input = CodeSearchToolInputSchema(
+        queries=["landslide nepal"], max_results=10
+    )
+
+    print("Running the search...")
+    output = await tool._arun(search_input)
+
+    print("\n--- Search Results ---")
+    for result in output.results:
+        print(result.url)
+        print(result.content)
+        print(result.extra["tool"])
+        print(result.extra["score"])
+        print("-" * 100)
+
+
 if __name__ == "__main__":
     print("Running local repo search test...")
     asyncio.run(local_repo_search_test())
     print("Running GitHub search test...")
     asyncio.run(github_search_test())
     print("Running SDE search test...")
     asyncio.run(sde_search_test())
+    print("Running combined code search test...")
+    asyncio.run(combined_code_search_test())

tests/code_search_tool_test.py

@@ -11,12 +11,11 @@
 
 import asyncio
 from akd.tools.misc import Embedder
-from akd.tools.search import SearxNGSearchToolConfig, SearxNGSearchToolInputSchema
+from akd.tools.search import SearxNGSearchToolConfig
 from akd.tools.code_search import (
     CodeSearchToolInputSchema,
     LocalRepoCodeSearchTool,
     LocalRepoCodeSearchToolConfig,
-    LocalRepoCodeSearchToolInputSchema,
     GitHubCodeSearchTool,
     SDECodeSearchTool,
     SDECodeSearchToolConfig,
@@ -99,7 +98,7 @@ def test_vector_embedding(embedder=Embedder(model_name="all-MiniLM-L6-v2")):
 
 @pytest.mark.asyncio
 async def test_local_repo_search(local_tool):
-    input_params = LocalRepoCodeSearchToolInputSchema(
+    input_params = CodeSearchToolInputSchema(
         queries=["landslide nepal"],
         max_results=3,
     )