Merge branch 'main' into huggingface

src/know_lang_bot/chat_bot/chat_graph.py

@@ -1,20 +1,85 @@
 # __future__ annotations is necessary for the type hints to work in this file
 from __future__ import annotations
 from dataclasses import dataclass
-from typing import List, Dict, Any, Optional
+from typing import AsyncGenerator, List, Dict, Any, Optional
 import chromadb
 from pydantic import BaseModel
-from pydantic_graph import BaseNode, Graph, GraphRunContext, End
+from pydantic_graph import BaseNode, EndStep, Graph, GraphRunContext, End, HistoryStep
 import ollama
 from know_lang_bot.config import AppConfig
 from know_lang_bot.utils.fancy_log import FancyLogger
 from pydantic_ai import Agent
 import logfire
 from pprint import pformat
+from enum import Enum
 
 LOG = FancyLogger(__name__)
 
-# Data Models
+class ChatStatus(str, Enum):
+    """Enum for tracking chat progress status"""
+    STARTING = "starting"
+    POLISHING = "polishing"
+    RETRIEVING = "retrieving"
+    ANSWERING = "answering"
+    COMPLETE = "complete"
+    ERROR = "error"
+
+class StreamingChatResult(BaseModel):
+    """Extended chat result with streaming information"""
+    answer: str
+    retrieved_context: Optional[RetrievedContext] = None
+    status: ChatStatus
+    progress_message: str
+    
+    @classmethod
+    def from_node(cls, node: BaseNode, state: ChatGraphState) -> StreamingChatResult:
+        """Create a StreamingChatResult from a node's current state"""
+        if isinstance(node, PolishQuestionNode):
+            return cls(
+                answer="",
+                status=ChatStatus.POLISHING,
+                progress_message=f"Refining question: '{state.original_question}'"
+            )
+        elif isinstance(node, RetrieveContextNode):
+            return cls(
+                answer="",
+                status=ChatStatus.RETRIEVING,
+                progress_message=f"Searching codebase with: '{state.polished_question or state.original_question}'"
+            )
+        elif isinstance(node, AnswerQuestionNode):
+            context_msg = f"Found {len(state.retrieved_context.chunks)} relevant segments" if state.retrieved_context else "No context found"
+            return cls(
+                answer="",
+                retrieved_context=state.retrieved_context,
+                status=ChatStatus.ANSWERING,
+                progress_message=f"Generating answer... {context_msg}"
+            )
+        else:
+            return cls(
+                answer="",
+                status=ChatStatus.ERROR,
+                progress_message=f"Unknown node type: {type(node).__name__}"
+            )
+    
+    @classmethod
+    def complete(cls, result: ChatResult) -> StreamingChatResult:
+        """Create a completed StreamingChatResult"""
+        return cls(
+            answer=result.answer,
+            retrieved_context=result.retrieved_context,
+            status=ChatStatus.COMPLETE,
+            progress_message="Response complete"
+        )
+    
+    @classmethod
+    def error(cls, error_msg: str) -> StreamingChatResult:
+        """Create an error StreamingChatResult"""
+        return cls(
+            answer=f"Error: {error_msg}",
+            status=ChatStatus.ERROR,
+            progress_message=f"An error occurred: {error_msg}"
+        )
+
 class RetrievedContext(BaseModel):
     """Structure for retrieved context"""
     chunks: List[str]
@@ -206,4 +271,66 @@ async def process_chat(
             answer="I encountered an error processing your question. Please try again."
         )
     finally:
-        return result
+        return result
+    
+async def stream_chat_progress(
+    question: str,
+    collection: chromadb.Collection,
+    config: AppConfig
+) -> AsyncGenerator[StreamingChatResult, None]:
+    """
+    Stream chat progress through the graph.
+    This is the main entry point for chat processing.
+    """
+    state = ChatGraphState(original_question=question)
+    deps = ChatGraphDeps(collection=collection, config=config)
+    
+    start_node = PolishQuestionNode()
+    history: list[HistoryStep[ChatGraphState, ChatResult]] = []
+
+    try:
+        # Initial status
+        yield StreamingChatResult(
+            answer="",
+            status=ChatStatus.STARTING,
+            progress_message=f"Processing question: {question}"
+        )
+
+        with logfire.span(
+            '{graph_name} run {start=}',
+            graph_name='RAG_chat_graph',
+            start=start_node,
+        ) as run_span:
+            current_node = start_node
+            
+            while True:
+                # Yield current node's status before processing
+                yield StreamingChatResult.from_node(current_node, state)
+                
+                try:
+                    # Process the current node
+                    next_node = await chat_graph.next(current_node, history, state=state, deps=deps, infer_name=False)
+                    
+                    if isinstance(next_node, End):
+                        result: ChatResult = next_node.data
+                        history.append(EndStep(result=next_node))
+                        run_span.set_attribute('history', history)
+                        # Yield final result
+                        yield StreamingChatResult.complete(result)
+                        return
+                    elif isinstance(next_node, BaseNode):
+                        current_node = next_node
+                    else:
+                        raise ValueError(f"Invalid node type: {type(next_node)}")
+                        
+                except Exception as node_error:
+                    LOG.error(f"Error in node {current_node.__class__.__name__}: {node_error}")
+                    yield StreamingChatResult.error(str(node_error))
+                    return
+                    
+    except Exception as e:
+        LOG.error(f"Error in stream_chat_progress: {e}")
+        yield StreamingChatResult.error(str(e))
+        return
+
+    

src/know_lang_bot/chat_bot/chat_interface.py

@@ -1,11 +1,12 @@
 import gradio as gr
 from know_lang_bot.config import AppConfig
 from know_lang_bot.utils.fancy_log import FancyLogger
-from know_lang_bot.chat_bot.chat_graph import ChatResult, process_chat
+from know_lang_bot.chat_bot.chat_graph import stream_chat_progress, ChatStatus
 import chromadb
-from typing import List, Dict
-import logfire
+from typing import List, Dict, AsyncGenerator
 from pathlib import Path
+from gradio import ChatMessage
+
 
 LOG = FancyLogger(__name__)
 
@@ -35,16 +36,86 @@ class CodeQAChatInterface:
         except Exception as e:
             LOG.error(f"Error reading code block: {e}")
             return "Error reading code"
-    
-    @logfire.instrument('Chatbot Process Question with {message=}')
-    async def process_question(
+
+    def _format_code_block(self, metadata: Dict) -> str:
+        """Format a single code block with metadata"""
+        file_path = metadata['file_path']
+        start_line = metadata['start_line']
+        end_line = metadata['end_line']
+        
+        code = self._get_code_block(file_path, start_line, end_line)
+        if not code:
+            return None
+            
+        title = f"📄 {file_path} (lines {start_line}-{end_line})"
+        if metadata.get('name'):
+            title += f" - {metadata['type']}: {metadata['name']}"
+            
+
+        return f"<details><summary>{title}</summary>\n\n```python\n{code}\n```\n\n</details>"
+
+    async def stream_response(
         self,
         message: str,
-        history: List[Dict[str, str]]
-    ) -> ChatResult:
-        """Process a question and return the answer with references"""
-        return await process_chat(message, self.collection, self.config)
-    
+        history: List[ChatMessage]
+    ) -> AsyncGenerator[List[ChatMessage], None]:
+        """Stream chat responses with progress updates"""
+        # Add user message
+        history.append(ChatMessage(role="user", content=message))
+        yield history
+        
+        current_progress: ChatMessage | None = None
+        code_blocks_added = False
+        
+        async for result in stream_chat_progress(message, self.collection, self.config):
+            # Handle progress updates
+            if result.status != ChatStatus.COMPLETE:
+                if current_progress:
+                    history.remove(current_progress)
+                
+                current_progress = ChatMessage(
+                    role="assistant",
+                    content=result.progress_message,
+                    metadata={
+                        "title": f"{result.status.value.title()} Progress",
+                        "status": "pending" if result.status != ChatStatus.ERROR else "error"
+                    }
+                )
+                history.append(current_progress)
+                yield history
+                continue
+
+            # When complete, remove progress message and add final content
+            if current_progress:
+                history.remove(current_progress)
+                current_progress = None
+
+            # Add code blocks before final answer if not added yet
+            if not code_blocks_added and result.retrieved_context and result.retrieved_context.metadatas:
+                total_code_blocks = []
+                for metadata in result.retrieved_context.metadatas:
+                    code_block = self._format_code_block(metadata)
+                    if code_block:
+                        total_code_blocks.append(code_block)
+
+                code_blocks_added = True
+                history.append(ChatMessage(
+                    role="assistant",
+                    content='\n\n'.join(total_code_blocks),
+                    metadata={
+                        "title": "💻 Code Context",
+                        "collapsible": True
+                    }
+                ))
+                yield history
+
+            # Add final answer
+            history.append(ChatMessage(
+                role="assistant",
+                content=result.answer
+            ))
+            yield history
+
     def create_interface(self) -> gr.Blocks:
         """Create the Gradio interface"""
         with gr.Blocks() as interface:
@@ -54,66 +125,28 @@ class CodeQAChatInterface:
             chatbot = gr.Chatbot(
                 type="messages",
                 bubble_full_width=False,
-                render_markdown=True
+                render_markdown=True,
+                height=600
             )
             
             msg = gr.Textbox(
                 label="Ask about the codebase",
                 placeholder="What does the CodeParser class do?",
-                container=False
+                container=False,
+                scale=7
             )
             
             with gr.Row():
-                submit = gr.Button("Submit")
-                clear = gr.ClearButton([msg, chatbot])
-
-            async def respond(message, history):
-                result = await self.process_question(message, history)
-                
-                # Format the answer with code blocks
-                formatted_messages = []
-                
-                # Add user message
-                formatted_messages.append({
-                    "role": "user", 
-                    "content": message
-                })
-                
-                # Collect code blocks first
-                code_blocks = []
-                if result.retrieved_context and result.retrieved_context.metadatas:
-                    for metadata in result.retrieved_context.metadatas:
-                        file_path = metadata['file_path']
-                        start_line = metadata['start_line']
-                        end_line = metadata['end_line']
-                        
-                        code = self._get_code_block(file_path, start_line, end_line)
-                        if code:
-                            title = f"📄 {file_path} (lines {start_line}-{end_line})"
-                            if metadata.get('name'):
-                                title += f" - {metadata['type']}: {metadata['name']}"
-                            
-                            code_blocks.append({
-                                "role": "assistant",
-                                "content": f"<details><summary>{title}</summary>\n\n```python\n{code}\n```\n\n</details>",
-                            })
-                
-                # Add code blocks before the answer
-                formatted_messages.extend(code_blocks)
-                
-                # Add assistant's answer
-                formatted_messages.append({
-                    "role": "assistant",
-                    "content": result.answer
-                })
-                
-                return {
-                    msg: "",
-                    chatbot: history + formatted_messages
-                }
+                submit = gr.Button("Submit", scale=1)
+                clear = gr.ClearButton([msg, chatbot], scale=1)
 
-            msg.submit(respond, [msg, chatbot], [msg, chatbot])
-            submit.click(respond, [msg, chatbot], [msg, chatbot])
+            async def respond(message: str, history: List[ChatMessage]) -> AsyncGenerator[List[ChatMessage], None]:
+                async for updated_history in self.stream_response(message, history):
+                    yield updated_history
+                    
+            # Set up event handlers
+            msg.submit(respond, [msg, chatbot], [chatbot])
+            submit.click(respond, [msg, chatbot], [chatbot])
 
         return interface