Make it possible to override the QA agent system prompt using the client

This commit is contained in:
Yiorgis Gozadinos 2025-10-07 13:15:58 +03:00
parent 496f2777cb
commit 0991df6d9e
No known key found for this signature in database
2 changed files with 18 additions and 2 deletions

View file

@ -207,6 +207,19 @@ answer = await client.ask("Who is the author of haiku.rag?", cite=True)
print(answer)
```
Customize the QA agent's behavior with a custom system prompt:
```python
custom_prompt = """You are a technical support expert for WIX.
Answer questions based on the knowledge base documents provided.
Be concise and helpful."""
answer = await client.ask(
"How do I create a blog?",
system_prompt=custom_prompt
)
```
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. Citations prefer the document title when present, otherwise they use the URI.
The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).

View file

@ -525,19 +525,22 @@ class HaikuRAG:
merged.append(current)
return merged
async def ask(self, question: str, cite: bool = False) -> str:
async def ask(
self, question: str, cite: bool = False, system_prompt: str | None = None
) -> str:
"""Ask a question using the configured QA agent.
Args:
question: The question to ask.
cite: Whether to include citations in the response.
system_prompt: Optional custom system prompt for the QA agent.
Returns:
The generated answer as a string.
"""
from haiku.rag.qa import get_qa_agent
qa_agent = get_qa_agent(self, use_citations=cite)
qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
return await qa_agent.answer(question)
async def rebuild_database(self) -> AsyncGenerator[str, None]: