From beb794b62e70d4e3b644179649caf44f43503f23 Mon Sep 17 00:00:00 2001 From: Yash Krishan Date: Wed, 31 Dec 2025 16:51:43 +0530 Subject: [PATCH 1/3] feat: enable chat before inference completes - Add INFERRING status to ProjectStatusEnum - Set project status to INFERRING at start of inference, READY on completion - Pass project_status to ChatContext for all agents - Exclude embedding-dependent tools (ask_knowledge_graph_queries, get_nodes_from_tags) during INFERRING - Add graceful error handling for vector index queries during INFERRING - Update all system agents and custom agents to respect INFERRING status - Add database migration for INFERRING status This allows users to chat with the AI agent even while AI enrichment (embeddings, docstrings) is still in progress, with limited tool access. --- .../versions/20251209_add_inferring_status.py | 41 +++++ .../conversation/conversation_service.py | 8 + app/modules/intelligence/agents/chat_agent.py | 6 + .../system_agents/blast_radius_agent.py | 13 +- .../system_agents/code_gen_agent.py | 18 +- .../chat_agents/system_agents/debug_agent.py | 17 +- .../system_agents/low_level_design_agent.py | 18 +- .../chat_agents/system_agents/qna_agent.py | 18 +- .../agents/custom_agents/runtime_agent.py | 21 ++- .../ask_knowledge_graph_queries_tool.py | 41 +++-- .../get_nodes_from_tags_tool.py | 22 ++- .../intelligence/tools/tool_service.py | 22 ++- .../knowledge_graph/inference_service.py | 160 +++++++++++------- app/modules/projects/projects_schema.py | 1 + 14 files changed, 292 insertions(+), 114 deletions(-) create mode 100644 app/alembic/versions/20251209_add_inferring_status.py diff --git a/app/alembic/versions/20251209_add_inferring_status.py b/app/alembic/versions/20251209_add_inferring_status.py new file mode 100644 index 00000000..dd759962 --- /dev/null +++ b/app/alembic/versions/20251209_add_inferring_status.py @@ -0,0 +1,41 @@ +"""Add inferring status to projects + +Revision ID: 20251209_add_inferring_status +Revises: 20251204181617_19b6f2ee95e6 +Create Date: 2025-12-09 16:45:00.000000 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "20251209_add_inferring_status" +down_revision: Union[str, None] = "20251204181617_19b6f2ee95e6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Drop the existing check constraint + op.drop_constraint("check_status", "projects", type_="check") + + # Create the new check constraint with 'inferring' status + op.create_check_constraint( + "check_status", + "projects", + "status IN ('submitted', 'cloned', 'parsed', 'inferring', 'ready', 'error')", + ) + + +def downgrade() -> None: + # Drop the constraint with 'inferring' + op.drop_constraint("check_status", "projects", type_="check") + + # Restore the original constraint without 'inferring' + op.create_check_constraint( + "check_status", + "projects", + "status IN ('submitted', 'cloned', 'parsed', 'ready', 'error')", + ) diff --git a/app/modules/conversations/conversation/conversation_service.py b/app/modules/conversations/conversation/conversation_service.py index 3057a017..9d6448f9 100644 --- a/app/modules/conversations/conversation/conversation_service.py +++ b/app/modules/conversations/conversation/conversation_service.py @@ -670,6 +670,12 @@ async def _generate_and_stream_ai_response( project_ids=[project_id] ) + # Get project status to conditionally enable/disable tools + project_info = await self.project_service.get_project_from_db_by_id( + project_id + ) + project_status = project_info.get("status") if project_info else None + # Prepare multimodal context - use current message attachments if available image_attachments = None if attachment_ids: @@ -702,6 +708,7 @@ async def _generate_and_stream_ai_response( history=validated_history[-12:], node_ids=[node.node_id for node in node_ids], query=query, + project_status=project_status, ), ) ) @@ -733,6 +740,7 @@ async def _generate_and_stream_ai_response( history=validated_history[-8:], node_ids=nodes, query=query, + project_status=project_status, image_attachments=image_attachments, context_images=context_images, ) diff --git a/app/modules/intelligence/agents/chat_agent.py b/app/modules/intelligence/agents/chat_agent.py index 045ea2bd..02b63c71 100644 --- a/app/modules/intelligence/agents/chat_agent.py +++ b/app/modules/intelligence/agents/chat_agent.py @@ -51,6 +51,8 @@ class ChatContext(BaseModel): node_ids: Optional[List[str]] = None additional_context: str = "" query: str + # Project parsing status - used to conditionally enable/disable tools + project_status: Optional[str] = None # Multimodal support - images attached to the current message image_attachments: Optional[Dict[str, Dict[str, Union[str, int]]]] = ( None # attachment_id -> {base64, mime_type, file_size, etc} @@ -58,6 +60,10 @@ class ChatContext(BaseModel): # Context images from recent conversation history context_images: Optional[Dict[str, Dict[str, Union[str, int]]]] = None + def is_inferring(self) -> bool: + """Check if the project is still in INFERRING state (AI enrichment in progress)""" + return self.project_status == "inferring" + def has_images(self) -> bool: """Check if this context contains any images""" return bool(self.image_attachments) or bool(self.context_images) diff --git a/app/modules/intelligence/agents/chat_agents/system_agents/blast_radius_agent.py b/app/modules/intelligence/agents/chat_agents/system_agents/blast_radius_agent.py index 7b1d722b..736d61d0 100644 --- a/app/modules/intelligence/agents/chat_agents/system_agents/blast_radius_agent.py +++ b/app/modules/intelligence/agents/chat_agents/system_agents/blast_radius_agent.py @@ -22,7 +22,7 @@ def __init__( self.llm_provider = llm_provider self.prompt_provider = prompt_provider - def _build_agent(self): + def _build_agent(self, ctx: ChatContext = None): agent_config = AgentConfig( role="Blast Radius Analyzer", goal="Analyze the impact of code changes", @@ -34,6 +34,10 @@ def _build_agent(self): ) ], ) + + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + tools = self.tools_provider.get_tools( [ "get_nodes_from_tags", @@ -46,7 +50,8 @@ def _build_agent(self): "fetch_file", "analyze_code_structure", "bash_command", - ] + ], + exclude_embedding_tools=exclude_embedding_tools, ) if not self.llm_provider.supports_pydantic("chat"): raise UnsupportedProviderError( @@ -55,12 +60,12 @@ def _build_agent(self): return PydanticRagAgent(self.llm_provider, agent_config, tools) async def run(self, ctx: ChatContext) -> ChatAgentResponse: - return await self._build_agent().run(ctx) + return await self._build_agent(ctx).run(ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: - async for chunk in self._build_agent().run_stream(ctx): + async for chunk in self._build_agent(ctx).run_stream(ctx): yield chunk diff --git a/app/modules/intelligence/agents/chat_agents/system_agents/code_gen_agent.py b/app/modules/intelligence/agents/chat_agents/system_agents/code_gen_agent.py index 6bbfbba4..f140414f 100644 --- a/app/modules/intelligence/agents/chat_agents/system_agents/code_gen_agent.py +++ b/app/modules/intelligence/agents/chat_agents/system_agents/code_gen_agent.py @@ -31,7 +31,7 @@ def __init__( self.tools_provider = tools_provider self.prompt_provider = prompt_provider - def _build_agent(self) -> ChatAgent: + def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: agent_config = AgentConfig( role="Code Generation Agent", goal="Generate precise, copy-paste ready code modifications that maintain project consistency and handle all dependencies", @@ -58,6 +58,14 @@ def _build_agent(self) -> ChatAgent: ) ], ) + + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + if exclude_embedding_tools: + logger.info( + "Project is in INFERRING status - excluding embedding-dependent tools" + ) + tools = self.tools_provider.get_tools( [ "get_code_from_multiple_node_ids", @@ -91,7 +99,8 @@ def _build_agent(self) -> ChatAgent: "fetch_file", "analyze_code_structure", "bash_command", - ] + ], + exclude_embedding_tools=exclude_embedding_tools, ) supports_pydantic = self.llm_provider.supports_pydantic("chat") should_use_multi = MultiAgentConfig.should_use_multi_agent( @@ -145,13 +154,14 @@ async def _enriched_context(self, ctx: ChatContext) -> ChatContext: return ctx async def run(self, ctx: ChatContext) -> ChatAgentResponse: - return await self._build_agent().run(await self._enriched_context(ctx)) + enriched_ctx = await self._enriched_context(ctx) + return await self._build_agent(enriched_ctx).run(enriched_ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: ctx = await self._enriched_context(ctx) - async for chunk in self._build_agent().run_stream(ctx): + async for chunk in self._build_agent(ctx).run_stream(ctx): yield chunk diff --git a/app/modules/intelligence/agents/chat_agents/system_agents/debug_agent.py b/app/modules/intelligence/agents/chat_agents/system_agents/debug_agent.py index d7ff9768..31a4276e 100644 --- a/app/modules/intelligence/agents/chat_agents/system_agents/debug_agent.py +++ b/app/modules/intelligence/agents/chat_agents/system_agents/debug_agent.py @@ -29,7 +29,7 @@ def __init__( self.llm_provider = llm_provider self.prompt_provider = prompt_provider - def _build_agent(self) -> ChatAgent: + def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: agent_config = AgentConfig( role="Context curation agent", goal="Handle querying the knowledge graph and refining the results to provide accurate and contextually rich responses.", @@ -49,6 +49,14 @@ def _build_agent(self) -> ChatAgent: ) ], ) + + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + if exclude_embedding_tools: + logger.info( + "Project is in INFERRING status - excluding embedding-dependent tools" + ) + tools = self.tools_provider.get_tools( [ "get_code_from_multiple_node_ids", @@ -82,7 +90,8 @@ def _build_agent(self) -> ChatAgent: "fetch_file", "analyze_code_structure", "bash_command", - ] + ], + exclude_embedding_tools=exclude_embedding_tools, ) supports_pydantic = self.llm_provider.supports_pydantic("chat") @@ -133,13 +142,13 @@ async def _enriched_context(self, ctx: ChatContext) -> ChatContext: async def run(self, ctx: ChatContext) -> ChatAgentResponse: ctx = await self._enriched_context(ctx) - return await self._build_agent().run(ctx) + return await self._build_agent(ctx).run(ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: ctx = await self._enriched_context(ctx) - async for chunk in self._build_agent().run_stream(ctx): + async for chunk in self._build_agent(ctx).run_stream(ctx): yield chunk diff --git a/app/modules/intelligence/agents/chat_agents/system_agents/low_level_design_agent.py b/app/modules/intelligence/agents/chat_agents/system_agents/low_level_design_agent.py index fc3d0827..78b67f19 100644 --- a/app/modules/intelligence/agents/chat_agents/system_agents/low_level_design_agent.py +++ b/app/modules/intelligence/agents/chat_agents/system_agents/low_level_design_agent.py @@ -29,7 +29,7 @@ def __init__( self.tools_provider = tools_provider self.prompt_provider = prompt_provider - def _build_agent(self) -> ChatAgent: + def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: agent_config = AgentConfig( role="Design Planner", goal="Create a detailed low-level design plan for implementing new features", @@ -45,6 +45,14 @@ def _build_agent(self) -> ChatAgent: ) ], ) + + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + if exclude_embedding_tools: + logger.info( + "Project is in INFERRING status - excluding embedding-dependent tools" + ) + tools = self.tools_provider.get_tools( [ "get_code_from_multiple_node_ids", @@ -77,7 +85,8 @@ def _build_agent(self) -> ChatAgent: "fetch_file", "analyze_code_structure", "bash_command", - ] + ], + exclude_embedding_tools=exclude_embedding_tools, ) supports_pydantic = self.llm_provider.supports_pydantic("chat") @@ -138,13 +147,14 @@ async def _enriched_context(self, ctx: ChatContext) -> ChatContext: return ctx async def run(self, ctx: ChatContext) -> ChatAgentResponse: - return await self._build_agent().run(await self._enriched_context(ctx)) + enriched_ctx = await self._enriched_context(ctx) + return await self._build_agent(enriched_ctx).run(enriched_ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: ctx = await self._enriched_context(ctx) - async for chunk in self._build_agent().run_stream(ctx): + async for chunk in self._build_agent(ctx).run_stream(ctx): yield chunk diff --git a/app/modules/intelligence/agents/chat_agents/system_agents/qna_agent.py b/app/modules/intelligence/agents/chat_agents/system_agents/qna_agent.py index 8da3ef2e..e6865fff 100644 --- a/app/modules/intelligence/agents/chat_agents/system_agents/qna_agent.py +++ b/app/modules/intelligence/agents/chat_agents/system_agents/qna_agent.py @@ -29,7 +29,7 @@ def __init__( self.tools_provider = tools_provider self.prompt_provider = prompt_provider - def _build_agent(self) -> ChatAgent: + def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: agent_config = AgentConfig( role="QNA Agent", goal="Answer queries of the repo in a detailed fashion", @@ -49,6 +49,14 @@ def _build_agent(self) -> ChatAgent: ) ], ) + + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + if exclude_embedding_tools: + logger.info( + "Project is in INFERRING status - excluding embedding-dependent tools (ask_knowledge_graph_queries, get_nodes_from_tags)" + ) + tools = self.tools_provider.get_tools( [ "get_code_from_multiple_node_ids", @@ -82,7 +90,8 @@ def _build_agent(self) -> ChatAgent: "fetch_file", "analyze_code_structure", "bash_command", - ] + ], + exclude_embedding_tools=exclude_embedding_tools, ) supports_pydantic = self.llm_provider.supports_pydantic("chat") @@ -143,13 +152,14 @@ async def _enriched_context(self, ctx: ChatContext) -> ChatContext: return ctx async def run(self, ctx: ChatContext) -> ChatAgentResponse: - return await self._build_agent().run(await self._enriched_context(ctx)) + enriched_ctx = await self._enriched_context(ctx) + return await self._build_agent(enriched_ctx).run(enriched_ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: ctx = await self._enriched_context(ctx) - async for chunk in self._build_agent().run_stream(ctx): + async for chunk in self._build_agent(ctx).run_stream(ctx): yield chunk diff --git a/app/modules/intelligence/agents/custom_agents/runtime_agent.py b/app/modules/intelligence/agents/custom_agents/runtime_agent.py index cc467ca2..739c1668 100644 --- a/app/modules/intelligence/agents/custom_agents/runtime_agent.py +++ b/app/modules/intelligence/agents/custom_agents/runtime_agent.py @@ -52,7 +52,7 @@ def __init__( self.tools_provider = tools_provider self.agent_config = CustomAgentConfig(**agent_config) - def _build_agent(self) -> ChatAgent: + def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: agent_config = AgentConfig( role=self.agent_config.role, goal=self.agent_config.goal, @@ -65,7 +65,17 @@ def _build_agent(self) -> ChatAgent: ], ) - tools = self.tools_provider.get_tools(self.agent_config.tasks[0].tools) + # Exclude embedding-dependent tools during INFERRING status + exclude_embedding_tools = ctx.is_inferring() if ctx else False + if exclude_embedding_tools: + logger.info( + "Project is in INFERRING status - excluding embedding-dependent tools for custom agent" + ) + + tools = self.tools_provider.get_tools( + self.agent_config.tasks[0].tools, + exclude_embedding_tools=exclude_embedding_tools + ) # Extract MCP servers from the first task with graceful error handling mcp_servers = [] @@ -133,13 +143,14 @@ async def _enriched_context(self, ctx: ChatContext) -> ChatContext: return ctx async def run(self, ctx: ChatContext) -> ChatAgentResponse: - return await self._build_agent().run(await self._enriched_context(ctx)) + enriched_ctx = await self._enriched_context(ctx) + return await self._build_agent(enriched_ctx).run(enriched_ctx) async def run_stream( self, ctx: ChatContext ) -> AsyncGenerator[ChatAgentResponse, None]: - ctx = await self._enriched_context(ctx) - async for chunk in self._build_agent().run_stream(ctx): + enriched_ctx = await self._enriched_context(ctx) + async for chunk in self._build_agent(enriched_ctx).run_stream(enriched_ctx): yield chunk diff --git a/app/modules/intelligence/tools/kg_based_tools/ask_knowledge_graph_queries_tool.py b/app/modules/intelligence/tools/kg_based_tools/ask_knowledge_graph_queries_tool.py index a4a0315b..9423d0e4 100644 --- a/app/modules/intelligence/tools/kg_based_tools/ask_knowledge_graph_queries_tool.py +++ b/app/modules/intelligence/tools/kg_based_tools/ask_knowledge_graph_queries_tool.py @@ -60,21 +60,34 @@ async def ask_multiple_knowledge_graph_queries( inference_service = InferenceService(self.sql_db, "dummy") async def process_query(query_request: QueryRequest) -> List[QueryResponse]: - # Call the query_vector_index method directly from InferenceService - results = inference_service.query_vector_index( - query_request.project_id, query_request.query, query_request.node_ids - ) - return [ - QueryResponse( - node_id=result.get("node_id"), - docstring=result.get("docstring"), - file_path=result.get("file_path"), - start_line=result.get("start_line") or 0, - end_line=result.get("end_line") or 0, - similarity=result.get("similarity"), + try: + # Call the query_vector_index method directly from InferenceService + results = inference_service.query_vector_index( + query_request.project_id, + query_request.query, + query_request.node_ids, + ) + return [ + QueryResponse( + node_id=result.get("node_id"), + docstring=result.get("docstring"), + file_path=result.get("file_path"), + start_line=result.get("start_line") or 0, + end_line=result.get("end_line") or 0, + similarity=result.get("similarity"), + ) + for result in results + ] + except Exception as e: + # Vector search may fail during INFERRING status (embeddings not ready) + # Return empty results gracefully instead of failing + import logging + + logging.warning( + f"Vector search failed for project {query_request.project_id} " + f"(likely during INFERRING): {e}" ) - for result in results - ] + return [] tasks = [process_query(query) for query in queries] results = await asyncio.gather(*tasks) diff --git a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py index 13bfe4b1..3ade7b4e 100644 --- a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py +++ b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py @@ -90,15 +90,21 @@ def run(self, tags: List[str], project_id: str) -> str: tag_conditions = " OR ".join([f"'{tag}' IN n.tags" for tag in tags]) query = f"""MATCH (n:NODE) WHERE ({tag_conditions}) AND n.repoId = '{project_id}' - RETURN n.file_path AS file_path, n.docstring AS docstring, n.text AS text, n.node_id AS node_id, n.name AS name + RETURN n.file_path AS file_path, COALESCE(n.docstring, substring(n.text, 0, 500)) AS docstring, n.text AS text, n.node_id AS node_id, n.name AS name """ - neo4j_config = ConfigProvider().get_neo4j_config() - nodes = CodeGraphService( - neo4j_config["uri"], - neo4j_config["username"], - neo4j_config["password"], - next(get_db()), - ).query_graph(query) + nodes = [] + try: + neo4j_config = ConfigProvider().get_neo4j_config() + nodes = CodeGraphService( + neo4j_config["uri"], + neo4j_config["username"], + neo4j_config["password"], + next(get_db()), + ).query_graph(query) + except Exception as e: + import logging + logging.warning(f"Error querying graph for tags for project {project_id}: {e}") + return [] return nodes diff --git a/app/modules/intelligence/tools/tool_service.py b/app/modules/intelligence/tools/tool_service.py index 52296a4a..ce1d94b1 100644 --- a/app/modules/intelligence/tools/tool_service.py +++ b/app/modules/intelligence/tools/tool_service.py @@ -97,6 +97,13 @@ class ToolService: + # Tools that depend on embeddings/docstrings and should be disabled during INFERRING status + # These tools require AI-inferenced content (embeddings, tags) that doesn't exist during inference + EMBEDDING_DEPENDENT_TOOLS = { + "ask_knowledge_graph_queries", # Uses vector search which requires embeddings + "get_nodes_from_tags", # Uses AI-generated tags which may not exist yet + } + def __init__(self, db: Session, user_id: str): self.db = db self.user_id = user_id @@ -123,10 +130,21 @@ def __init__(self, db: Session, user_id: str): self.provider_service = ProviderService.create(db, user_id) self.tools = self._initialize_tools() - def get_tools(self, tool_names: List[str]) -> List[StructuredTool]: - """get tools if exists""" + def get_tools( + self, tool_names: List[str], exclude_embedding_tools: bool = False + ) -> List[StructuredTool]: + """Get tools if they exist. + + Args: + tool_names: List of tool names to retrieve + exclude_embedding_tools: If True, excludes tools that depend on embeddings/AI-generated content. + Use this when project is in INFERRING status. + """ tools = [] for tool_name in tool_names: + # Skip embedding-dependent tools if requested (during INFERRING status) + if exclude_embedding_tools and tool_name in self.EMBEDDING_DEPENDENT_TOOLS: + continue if self.tools.get(tool_name) is not None: tools.append(self.tools[tool_name]) return tools diff --git a/app/modules/parsing/knowledge_graph/inference_service.py b/app/modules/parsing/knowledge_graph/inference_service.py index 86d84e14..c5455746 100644 --- a/app/modules/parsing/knowledge_graph/inference_service.py +++ b/app/modules/parsing/knowledge_graph/inference_service.py @@ -22,6 +22,7 @@ generate_content_hash, is_content_cacheable, ) +from app.modules.projects.projects_schema import ProjectStatusEnum from app.modules.projects.projects_service import ProjectService from app.modules.search.search_service import SearchService from app.modules.utils.logger import setup_logger @@ -290,7 +291,7 @@ def consolidate_chunk_responses( # Combine multiple chunk descriptions intelligently consolidated_text = f"This is a large code component split across {len(all_docstrings)} sections: " consolidated_text += " | ".join( - [f"Section {i+1}: {doc}" for i, doc in enumerate(all_docstrings)] + [f"Section {i + 1}: {doc}" for i, doc in enumerate(all_docstrings)] ) # Create single consolidated docstring for parent node @@ -1091,13 +1092,31 @@ def create_vector_index(self): ) async def run_inference(self, repo_id: str): - docstrings = await self.generate_docstrings(repo_id) - logger.info( - f"DEBUGNEO4J: After generate docstrings, Repo ID: {repo_id}, Docstrings: {len(docstrings)}" - ) - self.log_graph_stats(repo_id) + try: + # Set status to INFERRING at the beginning + await self.project_manager.update_project_status( + repo_id, ProjectStatusEnum.INFERRING + ) - self.create_vector_index() + docstrings = await self.generate_docstrings(repo_id) + logger.info( + f"DEBUGNEO4J: After generate docstrings, Repo ID: {repo_id}, Docstrings: {len(docstrings)}" + ) + self.log_graph_stats(repo_id) + + self.create_vector_index() + + # Set status to READY after successful completion + await self.project_manager.update_project_status( + repo_id, ProjectStatusEnum.READY + ) + except Exception as e: + logger.error(f"Inference failed for project {repo_id}: {e}") + # Set status to ERROR on failure + await self.project_manager.update_project_status( + repo_id, ProjectStatusEnum.ERROR + ) + raise def query_vector_index( self, @@ -1106,65 +1125,76 @@ def query_vector_index( node_ids: Optional[List[str]] = None, top_k: int = 5, ) -> List[Dict]: + """ + Query the vector index for similar nodes. + + Note: This may fail if called during INFERRING status when embeddings/index + are not yet ready. The calling tool (ask_knowledge_graph_queries) handles + these errors gracefully by returning empty results. + """ embedding = self.generate_embedding(query) with self.driver.session() as session: - if node_ids: - # Fetch context node IDs - result_neighbors = session.run( - """ - MATCH (n:NODE) - WHERE n.repoId = $project_id AND n.node_id IN $node_ids - CALL { - WITH n - MATCH (n)-[*1..4]-(neighbor:NODE) - RETURN COLLECT(DISTINCT neighbor.node_id) AS neighbor_ids - } - RETURN COLLECT(DISTINCT n.node_id) + REDUCE(acc = [], neighbor_ids IN COLLECT(neighbor_ids) | acc + neighbor_ids) AS context_node_ids - """, - project_id=project_id, - node_ids=node_ids, - ) - context_node_ids = result_neighbors.single()["context_node_ids"] - - # Use vector index and filter by context_node_ids - result = session.run( - """ - CALL db.index.vector.queryNodes('docstring_embedding', $initial_k, $embedding) - YIELD node, score - WHERE node.repoId = $project_id AND node.node_id IN $context_node_ids - RETURN node.node_id AS node_id, - node.docstring AS docstring, - node.file_path AS file_path, - node.start_line AS start_line, - node.end_line AS end_line, - score AS similarity - ORDER BY similarity DESC - LIMIT $top_k - """, - project_id=project_id, - embedding=embedding, - context_node_ids=context_node_ids, - initial_k=top_k * 10, # Adjust as needed - top_k=top_k, - ) - else: - result = session.run( - """ - CALL db.index.vector.queryNodes('docstring_embedding', $top_k, $embedding) - YIELD node, score - WHERE node.repoId = $project_id - RETURN node.node_id AS node_id, - node.docstring AS docstring, - node.file_path AS file_path, - node.start_line AS start_line, - node.end_line AS end_line, - score AS similarity - """, - project_id=project_id, - embedding=embedding, - top_k=top_k, - ) + try: + if node_ids: + # Fetch context node IDs + result_neighbors = session.run( + """ + MATCH (n:NODE) + WHERE n.repoId = $project_id AND n.node_id IN $node_ids + CALL { + WITH n + MATCH (n)-[*1..4]-(neighbor:NODE) + RETURN COLLECT(DISTINCT neighbor.node_id) AS neighbor_ids + } + RETURN COLLECT(DISTINCT n.node_id) + REDUCE(acc = [], neighbor_ids IN COLLECT(neighbor_ids) | acc + neighbor_ids) AS context_node_ids + """, + project_id=project_id, + node_ids=node_ids, + ) + context_node_ids = result_neighbors.single()["context_node_ids"] + + # Use vector index and filter by context_node_ids + result = session.run( + """ + CALL db.index.vector.queryNodes('docstring_embedding', $initial_k, $embedding) + YIELD node, score + WHERE node.repoId = $project_id AND node.node_id IN $context_node_ids + RETURN node.node_id AS node_id, + node.docstring AS docstring, + node.file_path AS file_path, + node.start_line AS start_line, + node.end_line AS end_line, + score AS similarity + ORDER BY similarity DESC + LIMIT $top_k + """, + project_id=project_id, + embedding=embedding, + context_node_ids=context_node_ids, + initial_k=top_k * 10, # Adjust as needed + top_k=top_k, + ) + else: + result = session.run( + """ + CALL db.index.vector.queryNodes('docstring_embedding', $top_k, $embedding) + YIELD node, score + WHERE node.repoId = $project_id + RETURN node.node_id AS node_id, + node.docstring AS docstring, + node.file_path AS file_path, + node.start_line AS start_line, + node.end_line AS end_line, + score AS similarity + """, + project_id=project_id, + embedding=embedding, + top_k=top_k, + ) - # Ensure all fields are included in the final output - return [dict(record) for record in result] + # Ensure all fields are included in the final output + return [dict(record) for record in result] + except Exception as e: + logger.warning(f"Error querying vector index for project {project_id}: {e}") + return [] diff --git a/app/modules/projects/projects_schema.py b/app/modules/projects/projects_schema.py index 58286d1d..dc356f84 100644 --- a/app/modules/projects/projects_schema.py +++ b/app/modules/projects/projects_schema.py @@ -8,6 +8,7 @@ class ProjectStatusEnum(str, Enum): CLONED = "cloned" PARSED = "parsed" PROCESSING = "processing" + INFERRING = "inferring" READY = "ready" ERROR = "error" From 9bf65b8cd73eebc4dedb397d646992768fde2f97 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 31 Dec 2025 11:43:21 +0000 Subject: [PATCH 2/3] chore: Auto-fix pre-commit issues --- .../intelligence/agents/custom_agents/runtime_agent.py | 2 +- .../tools/kg_based_tools/get_nodes_from_tags_tool.py | 5 ++++- app/modules/parsing/knowledge_graph/inference_service.py | 6 ++++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/app/modules/intelligence/agents/custom_agents/runtime_agent.py b/app/modules/intelligence/agents/custom_agents/runtime_agent.py index 739c1668..5532cf04 100644 --- a/app/modules/intelligence/agents/custom_agents/runtime_agent.py +++ b/app/modules/intelligence/agents/custom_agents/runtime_agent.py @@ -74,7 +74,7 @@ def _build_agent(self, ctx: ChatContext = None) -> ChatAgent: tools = self.tools_provider.get_tools( self.agent_config.tasks[0].tools, - exclude_embedding_tools=exclude_embedding_tools + exclude_embedding_tools=exclude_embedding_tools, ) # Extract MCP servers from the first task with graceful error handling diff --git a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py index 3ade7b4e..1dd19eba 100644 --- a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py +++ b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py @@ -103,7 +103,10 @@ def run(self, tags: List[str], project_id: str) -> str: ).query_graph(query) except Exception as e: import logging - logging.warning(f"Error querying graph for tags for project {project_id}: {e}") + + logging.warning( + f"Error querying graph for tags for project {project_id}: {e}" + ) return [] return nodes diff --git a/app/modules/parsing/knowledge_graph/inference_service.py b/app/modules/parsing/knowledge_graph/inference_service.py index c5455746..dff64f14 100644 --- a/app/modules/parsing/knowledge_graph/inference_service.py +++ b/app/modules/parsing/knowledge_graph/inference_service.py @@ -1127,7 +1127,7 @@ def query_vector_index( ) -> List[Dict]: """ Query the vector index for similar nodes. - + Note: This may fail if called during INFERRING status when embeddings/index are not yet ready. The calling tool (ask_knowledge_graph_queries) handles these errors gracefully by returning empty results. @@ -1196,5 +1196,7 @@ def query_vector_index( # Ensure all fields are included in the final output return [dict(record) for record in result] except Exception as e: - logger.warning(f"Error querying vector index for project {project_id}: {e}") + logger.warning( + f"Error querying vector index for project {project_id}: {e}" + ) return [] From 0d9c3fd5761758c6fb622e559709096ce2fb57a1 Mon Sep 17 00:00:00 2001 From: Yash Krishan Date: Wed, 31 Dec 2025 17:22:22 +0530 Subject: [PATCH 3/3] fix: add missing 'processing' status to migration and fix DB session leak - Add 'processing' status to check constraint in migration to match ProjectStatusEnum - Fix DB session leak in get_nodes_from_tags_tool by properly managing generator lifecycle - Ensure generator.close() is called in finally block to trigger cleanup --- .../versions/20251209_add_inferring_status.py | 6 ++++-- .../kg_based_tools/get_nodes_from_tags_tool.py | 15 ++++++++++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/app/alembic/versions/20251209_add_inferring_status.py b/app/alembic/versions/20251209_add_inferring_status.py index dd759962..8bf32785 100644 --- a/app/alembic/versions/20251209_add_inferring_status.py +++ b/app/alembic/versions/20251209_add_inferring_status.py @@ -22,10 +22,11 @@ def upgrade() -> None: op.drop_constraint("check_status", "projects", type_="check") # Create the new check constraint with 'inferring' status + # Note: 'processing' is included to match ProjectStatusEnum.PROCESSING op.create_check_constraint( "check_status", "projects", - "status IN ('submitted', 'cloned', 'parsed', 'inferring', 'ready', 'error')", + "status IN ('submitted', 'cloned', 'parsed', 'processing', 'inferring', 'ready', 'error')", ) @@ -34,8 +35,9 @@ def downgrade() -> None: op.drop_constraint("check_status", "projects", type_="check") # Restore the original constraint without 'inferring' + # Note: 'processing' is included to match ProjectStatusEnum.PROCESSING op.create_check_constraint( "check_status", "projects", - "status IN ('submitted', 'cloned', 'parsed', 'ready', 'error')", + "status IN ('submitted', 'cloned', 'parsed', 'processing', 'ready', 'error')", ) diff --git a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py index 1dd19eba..356eb480 100644 --- a/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py +++ b/app/modules/intelligence/tools/kg_based_tools/get_nodes_from_tags_tool.py @@ -93,13 +93,17 @@ def run(self, tags: List[str], project_id: str) -> str: RETURN n.file_path AS file_path, COALESCE(n.docstring, substring(n.text, 0, 500)) AS docstring, n.text AS text, n.node_id AS node_id, n.name AS name """ nodes = [] + # Properly manage the DB generator to ensure cleanup + gen = get_db() + db = None try: + db = next(gen) neo4j_config = ConfigProvider().get_neo4j_config() nodes = CodeGraphService( neo4j_config["uri"], neo4j_config["username"], neo4j_config["password"], - next(get_db()), + db, ).query_graph(query) except Exception as e: import logging @@ -108,6 +112,15 @@ def run(self, tags: List[str], project_id: str) -> str: f"Error querying graph for tags for project {project_id}: {e}" ) return [] + finally: + # Close the generator to trigger its finally block, which closes the DB session + if gen: + try: + gen.close() + except (GeneratorExit, StopIteration): + # GeneratorExit is expected when closing a generator + # StopIteration may occur if generator is already exhausted + pass return nodes