QA LLM, Ollama implementation
This commit is contained in:
parent
26900d9687
commit
095d532a7a
7 changed files with 115 additions and 2 deletions
|
|
@ -19,6 +19,9 @@ class AppConfig(BaseModel):
|
||||||
EMBEDDINGS_MODEL: str = "mxbai-embed-large"
|
EMBEDDINGS_MODEL: str = "mxbai-embed-large"
|
||||||
EMBEDDINGS_VECTOR_DIM: int = 1024
|
EMBEDDINGS_VECTOR_DIM: int = 1024
|
||||||
|
|
||||||
|
QA_PROVIDER: str = "ollama"
|
||||||
|
QA_MODEL: str = "qwen3"
|
||||||
|
|
||||||
CHUNK_SIZE: int = 256
|
CHUNK_SIZE: int = 256
|
||||||
CHUNK_OVERLAP: int = 32
|
CHUNK_OVERLAP: int = 32
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -49,7 +49,6 @@ class FileWatcher:
|
||||||
try:
|
try:
|
||||||
uri = file.as_uri()
|
uri = file.as_uri()
|
||||||
existing_doc = await self.client.get_document_by_uri(uri)
|
existing_doc = await self.client.get_document_by_uri(uri)
|
||||||
print(uri)
|
|
||||||
if existing_doc:
|
if existing_doc:
|
||||||
doc = await self.client.create_document_from_source(str(file))
|
doc = await self.client.create_document_from_source(str(file))
|
||||||
logger.info(f"Updated document {existing_doc.id} from {file}")
|
logger.info(f"Updated document {existing_doc.id} from {file}")
|
||||||
|
|
|
||||||
0
src/haiku/rag/qa/__init__.py
Normal file
0
src/haiku/rag/qa/__init__.py
Normal file
16
src/haiku/rag/qa/base.py
Normal file
16
src/haiku/rag/qa/base.py
Normal file
|
|
@ -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."
|
||||||
|
)
|
||||||
89
src/haiku/rag/qa/ollama.py
Normal file
89
src/haiku/rag/qa/ollama.py
Normal file
|
|
@ -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"]
|
||||||
7
src/haiku/rag/qa/prompts.py
Normal file
7
src/haiku/rag/qa/prompts.py
Normal file
|
|
@ -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."
|
||||||
|
"""
|
||||||
|
|
@ -325,7 +325,6 @@ class ChunkRepository(BaseRepository[Chunk]):
|
||||||
words = re.findall(r"\b\w+\b", query.lower())
|
words = re.findall(r"\b\w+\b", query.lower())
|
||||||
# Join with OR to find chunks containing any of the keywords
|
# Join with OR to find chunks containing any of the keywords
|
||||||
fts_query = " OR ".join(words) if words else query
|
fts_query = " OR ".join(words) if words else query
|
||||||
|
|
||||||
# Perform hybrid search using RRF (Reciprocal Rank Fusion)
|
# Perform hybrid search using RRF (Reciprocal Rank Fusion)
|
||||||
cursor.execute(
|
cursor.execute(
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue