danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Pass context for CLI/app

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-16 12:40:35 +02:00
parent d82e1f95c1
commit 18362f01c2
No known key found for this signature in database
6 changed files with 120 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
app/backend/main.py
View file

@ -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,
)

27
haiku_rag_slim/haiku/rag/app.py
View file

@ -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)

6
haiku_rag_slim/haiku/rag/chat/__init__.py
View file

@ -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()

19
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -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."""

70
haiku_rag_slim/haiku/rag/cli.py
View file

@ -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(

8
tests/test_app.py
View file

@ -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