feat: implement codeSearch functionality (#4629)

Author: sahil485

servers/fai-lambda/fai-code-indexing/src/handler.py

@@ -35,16 +35,23 @@ def handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
             repo_urls = validate_body_param_or_throw(body, "repoUrls", list[str])
             asyncio.run(setup_repos_for_domain(domain=domain, repo_urls=repo_urls))
 
-            message = f"Successfully indexed {len(repo_urls)} repositories"
+            response_body = {
+                "message": f"Successfully indexed {len(repo_urls)} repositories",
+                "timestamp": datetime.now(UTC).isoformat(),
+                "requestId": context.aws_request_id,
+            }
+
         elif event_type == "codeSearch":
-            asyncio.run(run_code_search_tool_call(domain=domain))
-            message = "Code search completed successfully"
+            question = validate_body_param_or_throw(body, "question")
+            session_id = body.get("sessionId")
+            result = asyncio.run(run_code_search_tool_call(domain=domain, question=question, session_id=session_id))
 
-        response_body = {
-            "message": message,
-            "timestamp": datetime.now(UTC).isoformat(),
-            "requestId": context.aws_request_id,
-        }
+            response_body = {
+                "message": "Code search completed successfully",
+                "timestamp": datetime.now(UTC).isoformat(),
+                "requestId": context.aws_request_id,
+                "answer": result.get("answer", None),
+            }
 
         return {
             "statusCode": 200,

servers/fai-lambda/fai-code-indexing/src/operations/search.py

@@ -1,10 +1,35 @@
 import logging
+import os
+from pathlib import Path
 from typing import Any
 
+from claude_agent_sdk import (
+    AssistantMessage,
+    ClaudeAgentOptions,
+    TextBlock,
+    ToolUseBlock,
+    query,
+)
+
 logger = logging.getLogger()
 
 
-async def run_code_search_tool_call(domain: str) -> dict[str, Any]:
+CLAUDE_ANSWER_CODE_QUESTION_USER_PROMPT = """
+You are given a question who's answer might lie in one of the code repositories and your job is to answer it.
+
+Process:
+- Analyze the question and determine which repository or repositories are relevant to the question.
+- Use the relevant repositories to answer the question to the best of your ability.
+
+Guidelines:
+- Keep answers as concise as possible while providing all the necessary details to unblock the user.
+- Be sure to provide a balance of code examples and text explanations to keep answers engaging and easy to follow.
+- Refrain from providing `full working examples` (only snippets/excerpts to support text steps)
+
+Question: {question}
+"""
+
+async def run_code_search_tool_call(domain: str, question: str, session_id: str | None = None) -> dict[str, Any]:
     """Run a code search tool call for a domain.
 
     Args:
@@ -15,9 +40,36 @@ async def run_code_search_tool_call(domain: str) -> dict[str, Any]:
     """
     logger.info(f"Running code search for domain: {domain}")
 
-    # TODO: Implement code search logic
+    efs_root = Path(os.environ.get("HOME", "/mnt/efs"))
+    domain_folder = efs_root / domain
+
+    answer = None
+
+    async for message in query(
+        prompt=CLAUDE_ANSWER_CODE_QUESTION_USER_PROMPT.format(question=question),
+        options=ClaudeAgentOptions(
+            cwd=str(domain_folder),
+            disallowed_tools=["Write", "Delete", "Rename"],
+            resume=session_id,
+            fork_session=True
+        )
+    ):
+        if hasattr(message, 'subtype') and message.subtype == 'init':
+            session_id = message.data.get('session_id')
+            print(f"Session started with ID: {session_id}")
+
+        if isinstance(message, AssistantMessage):
+            for content in message.content:
+                if isinstance(content, ToolUseBlock):
+                    print(f"Tool used: {content.name}")
+                    print(f"Tool input: {content.input}")
+                if isinstance(content, TextBlock):
+                    print(f"Text: {content.text}")
+                    answer = content.text
+
 
     return {
         "domain": domain,
         "status": "success",
+        "answer": answer,
     }