Fix: Connect resource extraction to main chat flow for @document_name mentions

[DevHammad0]

Problem Statement

The resource extraction functionality was implemented but completely disconnected from the main chat flow. When users typed queries like "Tell me about @report.pdf", the system would:

• Show autocomplete for resources ✅
• Extract @ mentions in code ✅
• But never actually inject document content into the AI prompt ❌

The _extract_resources() and _process_query() methods existed but were never called, causing the AI to receive raw queries without document context.

---

What Was Broken Before:

Original Problem:

The resource extraction code existed but was completely disconnected from the main chat flow. Here's what happened:

# BROKEN FLOW:
User types: "Tell me about @report.pdf"
    ↓
cli.py: user_input = "Tell me about @report.pdf"
    ↓
cli.py: response = await self.agent.run(user_input)
    ↓
chat.py: async def run(self, query: str) -> str:
    ↓
chat.py: response = await self.agent_serve.chat(query=query, mcp_clients=self.clients)
    ↓
AI gets: "Tell me about @report.pdf" (RAW - no document content!)
    ↓
AI tries to use tools to read document with ID "@report.pdf" (including @)
    ↓
Tool responds: "Document with id @report.pdf not found" (@ is not part of actual doc ID)
    ↓
AI responds: "I am sorry, I cannot find the document `@report.pdf`."

The _extract_resources() and _process_query() methods existed in cli_chat.py but were NEVER CALLED!

Root Cause:

class CliChat(Chat):
    # Had _extract_resources() and _process_query() methods
    # But inherited Chat.run() which bypassed them completely!
    
    async def _extract_resources(self, query: str) -> str:
        # This method existed but was never called!
        pass
    
    async def _process_query(self, query: str):
        # This method existed but was never called!
        pass

Instead of using resources, the system was incorrectly trying to use tools with the "@" symbol included in the document ID, which doesn't exist.

---

What I Changed to Fix It:

Key Fix: Override the run() method in CliChat

After (working):

class CliChat(Chat):
    async def run(self, query: str) -> str:  # ← ADDED THIS METHOD
        """Override run method to process resources before sending to agent"""
        # Process query to extract resources and inject document content
        await self._process_query(query)
        
        # Call agent service without passing query since _process_query already added to messages
        response = await self.agent_serve.chat(
            query="",  # Empty because _process_query already added enhanced prompt to messages
            mcp_clients=self.clients,
        )
        
        return response.final_output

---

Current Working Flow:

Complete Step-by-Step Flow:

# WORKING FLOW:
User types: "Tell me about @report.pdf"
    ↓
cli.py: user_input = "Tell me about @report.pdf"
    ↓
cli.py: response = await self.agent.run(user_input)  # self.agent is CliChat
    ↓
cli_chat.py: async def run(self, query: str) -> str:  # ← NOW CALLS OUR OVERRIDE!
    ↓
cli_chat.py: await self._process_query(query)
    ↓
cli_chat.py: _process_query() calls _extract_resources(query)
    ↓
cli_chat.py: _extract_resources() does:
    - mentions = ["report.pdf"]  # Extract @mentions (removes @ symbol)
    - doc_ids = await self.list_docs_ids()  # Get all available docs
    - content = await self.get_doc_content("report.pdf")  # Fetch content using correct ID
    - Returns: '<document id="report.pdf">The report details...</document>'
    ↓
cli_chat.py: _process_query() creates enhanced prompt:
    """
    The user has a question:
    <query>Tell me about @report.pdf</query>
    
    The following context may be useful:
    <context>
    <document id="report.pdf">The report details the state of a 20m condenser tower.</document>
    </context>
    """
    ↓
cli_chat.py: self.agent_serve.messages.append({"role": "user", "content": enhanced_prompt})
    ↓
cli_chat.py: response = await self.agent_serve.chat(query="", mcp_clients=self.clients)
    ↓
AI gets: Enhanced prompt with document content already included!
    ↓
AI responds: "The report details the state of a 20m condenser tower."

---

Detailed Code Changes:

1. Added run() method override in CliChat (core/cli_chat.py):

async def run(self, query: str) -> str:
    """Override run method to process resources before sending to agent"""
    # Process query to extract resources and inject document content
    await self._process_query(query)
    
    # Call agent service without passing query since _process_query already added to messages
    response = await self.agent_serve.chat(
        query="",  # Empty because _process_query already added enhanced prompt to messages
        mcp_clients=self.clients,
    )
    
    return response.final_output

---

Why It Works Now:

Method Inheritance Chain:

# Before (broken):
CliApp.run() → CliChat.run() → Chat.run() → AgentService.chat()
                     ↑                ↑
                Didn't exist!    Bypassed resource processing!

# After (working):
CliApp.run() → CliChat.run() → CliChat._process_query() → AgentService.chat()
                     ↑                        ↑
                 Our override!        Processes resources!

Key Insight:

The original code had all the right pieces (_extract_resources, _process_query) but they were like disconnected functions that never got called. By overriding the run() method in CliChat, I connected the resource processing pipeline to the main chat flow.

Most importantly, this fix ensures that:

• Resources are used instead of tools for document access
• @ symbol is properly stripped from document IDs before resource lookup
• Document content is pre-injected into the prompt instead of requiring tool calls

---

Testing Results:

Before Fix:

$ uv run main.py
> Tell me about @report.pdf
Response: I am sorry, I cannot find the document `@report.pdf`.

(System was trying to call tool with ID "@report.pdf" which doesn't exist)

After Fix:

$ uv run main.py
> Tell me about @report.pdf
Response: The report details the state of a 20m condenser tower.

Multiple Resources:

$ uv run main.py
> Compare @deposition.md and @plan.md
Response: The deposition covers Angela Smith's testimony while the plan outlines implementation steps.

---

Impact:

✅ Before: AI tried to use tools with "@report.pdf" → tool failed → "document not found"
✅ After: AI gets enhanced prompt with document content via resources → answers perfectly

What Now Works:

• Resource Extraction: @document_name mentions are properly extracted and @ is stripped
• Content Injection: Document content is injected into AI prompts using resources
• Autocomplete: Still works for selecting resources
• Multiple Resources: Can handle multiple @ mentions in single query
• Error Handling: Graceful handling of missing documents
• Proper Resource Usage: Uses MCP resources instead of incorrectly trying to use tools

No Breaking Changes:

• All existing functionality preserved
• Autocomplete still works
• Tool functionality unaffected
• Prompt functionality unaffected

---

Files Modified:

1. core/cli_chat.py:
   • Added run() method override to connect resource processing

---

The fix was essentially connecting the dots between the existing resource extraction code and the main chat execution flow by overriding the run() method in CliChat. This ensures that the system properly uses MCP resources (not tools) for document access, making the functionality work as intended for the learn-agentic-ai course materials.