From 18362f01c22bbde6382acde221e6891d55316adc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 16 Jan 2026 12:40:35 +0200 Subject: [PATCH] Pass context for CLI/app --- app/backend/main.py | 5 +- haiku_rag_slim/haiku/rag/app.py | 27 +++++++-- haiku_rag_slim/haiku/rag/chat/__init__.py | 6 +- haiku_rag_slim/haiku/rag/chat/app.py | 19 ++++-- haiku_rag_slim/haiku/rag/cli.py | 70 ++++++++++++++++++++++- tests/test_app.py | 8 ++- 6 files changed, 120 insertions(+), 15 deletions(-) diff --git a/app/backend/main.py b/app/backend/main.py index 42a4fa5d..6479f0ef 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -80,8 +80,9 @@ async def stream_chat(request: Request) -> Response: accept = request.headers.get("accept", SSE_CONTENT_TYPE) run_input = AGUIAdapter.build_run_input(body) - # Restore qa_history from incoming state (look under namespaced key) + # Restore session state from incoming AG-UI state (look under namespaced key) initial_qa_history: list[QAResponse] = [] + initial_context: str | None = None state = getattr(run_input, "state", None) if state: chat_state = state.get(AGUI_STATE_KEY, state) @@ -89,6 +90,7 @@ async def stream_chat(request: Request) -> Response: initial_qa_history = [ QAResponse(**qa) for qa in chat_state.get("qa_history", []) ] + initial_context = chat_state.get("initial_context") # Build deps with session state thread_id = getattr(run_input, "thread_id", None) @@ -98,6 +100,7 @@ async def stream_chat(request: Request) -> Response: session_state=ChatSessionState( session_id=thread_id or "", qa_history=initial_qa_history, + initial_context=initial_context, ), state_key=AGUI_STATE_KEY, ) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 6305bc70..1167b61f 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -376,6 +376,7 @@ class HaikuRAGApp: cite: bool = False, deep: bool = False, filter: str | None = None, + initial_context: str | None = None, ): """Ask a question using the RAG system. @@ -384,6 +385,7 @@ class HaikuRAGApp: cite: Include citations in the answer deep: Use deep QA mode (multi-step reasoning) filter: SQL WHERE clause to filter documents + initial_context: Optional background context for the question """ async with HaikuRAG( db_path=self.db_path, @@ -394,7 +396,9 @@ class HaikuRAGApp: citations = [] if deep: graph = build_research_graph(config=self.config) - context = ResearchContext(original_question=question) + context = ResearchContext( + original_question=question, initial_context=initial_context + ) state = ResearchState.from_config( context=context, config=self.config, @@ -423,7 +427,14 @@ class HaikuRAGApp: else: self.console.print("[yellow]No answer generated.[/yellow]") else: - answer, citations = await self.client.ask(question, filter=filter) + system_prompt = ( + f"BACKGROUND CONTEXT:\n{initial_context}" + if initial_context + else None + ) + answer, citations = await self.client.ask( + question, system_prompt=system_prompt, filter=filter + ) self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() @@ -433,12 +444,18 @@ class HaikuRAGApp: for renderable in format_citations_rich(citations): self.console.print(renderable) - async def research(self, question: str, filter: str | None = None): + async def research( + self, + question: str, + filter: str | None = None, + initial_context: str | None = None, + ): """Run research via the pydantic-graph pipeline. Args: question: The research question filter: SQL WHERE clause to filter documents + initial_context: Optional background context for the research """ async with HaikuRAG( db_path=self.db_path, @@ -451,7 +468,9 @@ class HaikuRAGApp: self.console.print() graph = build_research_graph(config=self.config) - context = ResearchContext(original_question=question) + context = ResearchContext( + original_question=question, initial_context=initial_context + ) state = ResearchState.from_config(context=context, config=self.config) state.search_filter = filter deps = ResearchDeps(client=client) diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index 6045a8b2..38c5a297 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -6,6 +6,7 @@ def run_chat( db_path: Path | None = None, read_only: bool = False, before: datetime | None = None, + initial_context: str | None = None, ) -> None: """Run the chat TUI. @@ -13,6 +14,7 @@ def run_chat( db_path: Path to the LanceDB database. If None, uses default from config. read_only: Whether to open the database in read-only mode. before: Query database as it existed before this datetime. + initial_context: Optional background context for the conversation. """ try: from haiku.rag.chat.app import ChatApp @@ -27,5 +29,7 @@ def run_chat( if db_path is None: db_path = config.storage.data_dir / "haiku.rag.lancedb" - app = ChatApp(db_path, read_only=read_only, before=before) + app = ChatApp( + db_path, read_only=read_only, before=before, initial_context=initial_context + ) app.run() diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 94b4f68a..c2dbf1cc 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -85,12 +85,17 @@ class ChatApp(App): # type: ignore[misc] ] def __init__( - self, db_path: Path, read_only: bool = False, before: datetime | None = None + self, + db_path: Path, + read_only: bool = False, + before: datetime | None = None, + initial_context: str | None = None, ) -> None: super().__init__() self.db_path = db_path self.read_only = read_only self.before = before + self.initial_context = initial_context self.client: HaikuRAG | None = None self.config = get_config() self.agent: Agent[ChatDeps, str] | None = None @@ -121,7 +126,10 @@ class ChatApp(App): # type: ignore[misc] # Create agent and session state self.agent = create_chat_agent(self.config) - self.session_state = ChatSessionState(session_id=str(uuid.uuid4())) + self.session_state = ChatSessionState( + session_id=str(uuid.uuid4()), + initial_context=self.initial_context, + ) # Focus the input field self.query_one(Input).focus() @@ -266,8 +274,11 @@ class ChatApp(App): # type: ignore[misc] self._last_citations.clear() self._selected_citation_idx = None self._message_history.clear() - # Reset session state for fresh conversation - self.session_state = ChatSessionState(session_id=str(uuid.uuid4())) + # Reset session state for fresh conversation (preserve initial_context) + self.session_state = ChatSessionState( + session_id=str(uuid.uuid4()), + initial_context=self.initial_context, + ) def action_focus_input(self) -> None: """Focus the input field, or cancel if processing.""" diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 11649c17..2cc84dfc 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -338,9 +338,34 @@ def ask( "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + context: str | None = typer.Option( + None, + "--context", + help="Background context for the question", + ), + context_file: Path | None = typer.Option( + None, + "--context-file", + help="Path to a file containing background context", + ), ): + # Resolve initial context from flag or file + initial_context: str | None = None + if context_file: + initial_context = context_file.read_text() + elif context: + initial_context = context + app = create_app(db) - asyncio.run(app.ask(question=question, cite=cite, deep=deep, filter=filter)) + asyncio.run( + app.ask( + question=question, + cite=cite, + deep=deep, + filter=filter, + initial_context=initial_context, + ) + ) @cli.command("research", help="Run multi-agent research and output a concise report") @@ -357,9 +382,28 @@ def research( "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + context: str | None = typer.Option( + None, + "--context", + help="Background context for the research", + ), + context_file: Path | None = typer.Option( + None, + "--context-file", + help="Path to a file containing background context", + ), ): + # Resolve initial context from flag or file + initial_context: str | None = None + if context_file: + initial_context = context_file.read_text() + elif context: + initial_context = context + app = create_app(db) - asyncio.run(app.research(question=question, filter=filter)) + asyncio.run( + app.research(question=question, filter=filter, initial_context=initial_context) + ) @cli.command("settings", help="Display current configuration settings") @@ -547,12 +591,32 @@ def chat( "--db", help="Path to the LanceDB database file", ), + context: str | None = typer.Option( + None, + "--context", + help="Initial context/background information for the conversation", + ), + context_file: Path | None = typer.Option( + None, + "--context-file", + help="Path to a file containing initial context", + ), ): """Launch the chat TUI for conversational RAG.""" from haiku.rag.chat import run_chat db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb" - run_chat(db_path, read_only=_read_only, before=_before) + + # Resolve initial context from flag or file + initial_context: str | None = None + if context_file: + initial_context = context_file.read_text() + elif context: + initial_context = context + + run_chat( + db_path, read_only=_read_only, before=_before, initial_context=initial_context + ) @cli.command( diff --git a/tests/test_app.py b/tests/test_app.py index 8443ba36..7e794e6a 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -304,7 +304,9 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.ask("test question") - mock_client.ask.assert_called_once_with("test question", filter=None) + mock_client.ask.assert_called_once_with( + "test question", system_prompt=None, filter=None + ) @pytest.mark.asyncio @@ -333,7 +335,9 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.ask("test question", cite=True) - mock_client.ask.assert_called_once_with("test question", filter=None) + mock_client.ask.assert_called_once_with( + "test question", system_prompt=None, filter=None + ) # Verify print was called (once for answer, once for citations) assert mock_print.call_count >= 1