added llm and rag queries

kg_ocr/rag/__init__.py

@@ -0,0 +1,3 @@
+from .query import retrieve, ask_wllm
+
+__all__ = ["retrieve", "ask_wllm"]

kg_ocr/rag/query.py

@@ -0,0 +1,32 @@
+from txtai.embeddings import Embeddings
+from txtai import LLM
+import litellm
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+def retrieve(embeddings: Embeddings, query: str, limit: int = 3) -> list[dict]:
+    """Search embeddings and return results with scores"""
+    return embeddings.search(query, limit)
+
+def ask_wllm(embeddings: Embeddings, question: str, model: str = "openrouter/minimax/minimax-m2.5:free", limit: int = 3) -> str:
+    """RAG: retrieve context from embeddings, then answer with an LLM."""
+    results = retrieve(embeddings, question, limit)
+    context = "\n\n".join([r["text"] for r in results])
+
+    response = litellm.completion(
+        model=model,
+        messages=[
+            {
+                "role": "system",
+                "content": "Answer ONLY using the provided context. Cite which parts you're drawing from. If the context doesn't cover something, say 'not in my documents'."
+            },
+            {
+                "role": "user",
+                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}"
+            }
+        ]
+    )
+
+    return response.choices[0].message.content