feat: Phase 5 - citations, source highlighting, and UI polish

src/companion/orchestrator.py

@@ -12,7 +12,7 @@ import httpx
 from companion.config import Config
 from companion.memory import Message, SessionMemory
 from companion.prompts import build_system_prompt, format_conversation_history
-from companion.rag.search import SearchEngine
+from companion.rag.search import SearchEngine, SearchResult
 
 
 class ChatOrchestrator:
@@ -47,7 +47,7 @@ class ChatOrchestrator:
             Streaming response chunks (SSE format)
         """
         # Retrieve relevant context from RAG if enabled
-        retrieved_context: list[dict[str, Any]] = []
+        retrieved_context: list[SearchResult] = []
         if use_rag:
             try:
                 retrieved_context = self.search.search(
@@ -103,9 +103,21 @@ class ChatOrchestrator:
                 session_id,
                 "assistant",
                 full_response,
-                metadata={"rag_context_used": len(retrieved_context) > 0},
+                metadata={
+                    "rag_context_used": len(retrieved_context) > 0,
+                    "citations": [ctx.to_dict() for ctx in retrieved_context[:5]]
+                    if retrieved_context
+                    else [],
+                },
             )
 
+        # Yield citations after the main response
+        if retrieved_context:
+            citations_data = {
+                "citations": [ctx.to_dict() for ctx in retrieved_context[:5]]
+            }
+            yield f"data: {json.dumps(citations_data)}\n\n"
+
     async def _stream_llm_response(
         self, messages: list[dict[str, Any]]
     ) -> AsyncGenerator[str, None]: