diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 26a29dfe..270c0aea 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,6 +19,9 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 + QA_PROVIDER: str = "ollama" + QA_MODEL: str = "qwen3" + CHUNK_SIZE: int = 256 CHUNK_OVERLAP: int = 32 diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py index 97809e9e..618bb32a 100644 --- a/src/haiku/rag/monitor.py +++ b/src/haiku/rag/monitor.py @@ -49,7 +49,6 @@ class FileWatcher: try: uri = file.as_uri() existing_doc = await self.client.get_document_by_uri(uri) - print(uri) if existing_doc: doc = await self.client.create_document_from_source(str(file)) logger.info(f"Updated document {existing_doc.id} from {file}") diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py new file mode 100644 index 00000000..6c8f8359 --- /dev/null +++ b/src/haiku/rag/qa/base.py @@ -0,0 +1,16 @@ +from haiku.rag.client import HaikuRAG +from haiku.rag.qa.prompts import SYSTEM_PROMPT + + +class QABase: + _model: str = "" + _system_prompt: str = SYSTEM_PROMPT + + def __init__(self, client: HaikuRAG, model: str = ""): + self._model = model + self._client = client + + async def answer(self, question: str) -> str: + raise NotImplementedError( + "QABase is an abstract class. Please implement the answer method in a subclass." + ) diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py new file mode 100644 index 00000000..273317b1 --- /dev/null +++ b/src/haiku/rag/qa/ollama.py @@ -0,0 +1,89 @@ +from ollama import AsyncClient + +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.qa.base import QABase + + +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={"temperature": 0.0, "seed": 42}, + 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={"temperature": 0.0, "seed": 42}, + ) + return final_response["message"]["content"] + else: + return response["message"]["content"] diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py new file mode 100644 index 00000000..fc8f2c9b --- /dev/null +++ b/src/haiku/rag/qa/prompts.py @@ -0,0 +1,7 @@ +SYSTEM_PROMPT = """ +You are a helpful assistant that uses a RAG library to answer the user's prompt. +Your task is to provide a concise and accurate answer based on the provided context. +You should ask the provided tools to find relevant documents and then use the content of those documents to answer the question. +Never make up information, always use the context to answer the question. +If the context does not contain enough information to answer the question, respond with "I cannot answer that based on the provided context." +""" diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index b80e27db..1ffec47b 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -325,7 +325,6 @@ class ChunkRepository(BaseRepository[Chunk]): words = re.findall(r"\b\w+\b", query.lower()) # Join with OR to find chunks containing any of the keywords fts_query = " OR ".join(words) if words else query - # Perform hybrid search using RRF (Reciprocal Rank Fusion) cursor.execute( """