91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
from ollama import AsyncClient
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.qa.base import QABase
|
|
|
|
OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 64000}
|
|
|
|
|
|
class QA(QABase):
|
|
def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL):
|
|
super().__init__(client, model or self._model)
|
|
|
|
async def answer(self, question: str) -> str:
|
|
ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
|
|
|
|
# Define the search tool
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_documents",
|
|
"description": "Search the knowledge base for relevant documents",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query to find relevant documents",
|
|
},
|
|
"limit": {
|
|
"type": "integer",
|
|
"description": "Maximum number of results to return",
|
|
"default": 3,
|
|
},
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
messages = [
|
|
{"role": "system", "content": self._system_prompt},
|
|
{"role": "user", "content": question},
|
|
]
|
|
|
|
# Initial response with tool calling
|
|
response = await ollama_client.chat(
|
|
model=self._model,
|
|
messages=messages,
|
|
tools=tools,
|
|
options=OLLAMA_OPTIONS,
|
|
think=False,
|
|
)
|
|
|
|
if response.get("message", {}).get("tool_calls"):
|
|
for tool_call in response["message"]["tool_calls"]:
|
|
if tool_call["function"]["name"] == "search_documents":
|
|
args = tool_call["function"]["arguments"]
|
|
query = args.get("query", question)
|
|
limit = int(args.get("limit", 3))
|
|
|
|
search_results = await self._client.search(query, limit=limit)
|
|
|
|
context_chunks = []
|
|
for chunk, score in search_results:
|
|
context_chunks.append(
|
|
f"Content: {chunk.content}\nScore: {score:.4f}"
|
|
)
|
|
|
|
context = "\n\n".join(context_chunks)
|
|
|
|
messages.append(response["message"])
|
|
messages.append(
|
|
{
|
|
"role": "tool",
|
|
"content": context,
|
|
"tool_call_id": tool_call.get("id", "search_tool"),
|
|
}
|
|
)
|
|
|
|
final_response = await ollama_client.chat(
|
|
model=self._model,
|
|
messages=messages,
|
|
think=False,
|
|
options=OLLAMA_OPTIONS,
|
|
)
|
|
return final_response["message"]["content"]
|
|
else:
|
|
return response["message"]["content"]
|