kg: graph build, traversal queries, neo4j export - kg_ocr.graph builds a networkx graph (docs, chunks, entities, citations, co-occurrence) from chunk markdown - analyzer for summaries, top entities/citations, anomaly checks - traversal: chunks_for_entity/citation, related_entities, expand_context - export: JSON round-trip, GraphML, batched MERGE into neo4j - new CLI: ocr-pipeline kg build|stats|query|export - lazy kg_ocr imports, networkx/neo4j behind extras - dropped dead watch.py shim, added KgConfig stub - trimmed README, updated TODO

kg_ocr/rag/__init__.py

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

kg_ocr/rag/query.py

@@ -1,32 +1,47 @@
-from txtai.embeddings import Embeddings
-from txtai import LLM
-import litellm
-from dotenv import load_dotenv
-import os
+from __future__ import annotations
 
-load_dotenv()
+from typing import Any
 
-def retrieve(embeddings: Embeddings, query: str, limit: int = 3) -> list[dict]:
-    """Search embeddings and return results with scores"""
+DEFAULT_CHAT_MODEL = "openrouter/minimax/minimax-m2.5:free"
+
+SYSTEM_PROMPT = (
+    "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'."
+)
+
+
+def retrieve(embeddings: Any, 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:
+
+def ask_wllm(
+    embeddings: Any,
+    question: str,
+    model: str = DEFAULT_CHAT_MODEL,
+    limit: int = 3,
+) -> str:
     """RAG: retrieve context from embeddings, then answer with an LLM."""
+    try:
+        import litellm
+    except ImportError as exc:
+        raise ImportError("ask_wllm needs litellm: uv sync --extra kg") from exc
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
     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": "system", "content": SYSTEM_PROMPT},
             {
                 "role": "user",
-                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}"
-            }
-        ]
+                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}",
+            },
+        ],
     )
 
-    return response.choices[0].message.content
+    return response.choices[0].message.content