File size: 4,059 Bytes
58b74a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """
knowledge_api.py — Public interface for the Knowledge Agent
============================================================
This module is the integration layer for the Master Orchestrator and any
other agent that needs to query your personal document store.
Usage from another agent:
from knowledge_agent.knowledge_api import query_knowledge, index_docs
It is intentionally stateless at the call level — the KnowledgeIndexer
caches the ChromaDB connection internally for performance.
"""
from typing import List, Dict, Optional
from indexer import KnowledgeIndexer
from llm import generate_answer, stream_answer
# Module-level singleton — one DB connection shared across all calls
# (safe for single-process use; FastAPI uses one process by default)
_indexer: Optional[KnowledgeIndexer] = None
def _get_indexer() -> KnowledgeIndexer:
global _indexer
if _indexer is None:
_indexer = KnowledgeIndexer()
return _indexer
# ── Core public functions ─────────────────────────────────────────────────────
def query_knowledge(
question: str,
top_k: int = 5,
return_sources: bool = True,
) -> Dict:
"""
PRIMARY ENTRY POINT for the Master Orchestrator and other agents.
Ask a question against your indexed personal documents.
Args:
question: Natural language question.
top_k: Number of chunks to retrieve from ChromaDB.
return_sources: If True, include source filenames in the response.
Returns:
{
"answer": str, # LLM-generated, grounded answer
"sources": [str, ...], # source filenames (if return_sources=True)
"chunks_used": int, # number of relevant chunks found
"question": str, # echo back for logging
}
Example:
>>> result = query_knowledge("What is our refund policy?")
>>> print(result["answer"])
"""
indexer = _get_indexer()
chunks = indexer.query(question, top_k=top_k)
if not chunks:
return {
"answer": "No documents have been indexed yet. "
"Please run `index_docs()` first.",
"sources": [],
"chunks_used": 0,
"question": question,
}
result = generate_answer(question, chunks)
result["question"] = question
if not return_sources:
result.pop("sources", None)
return result
def index_docs(docs_dir: str = "./documents", force: bool = False):
"""
Index (or re-index) all supported documents in docs_dir.
Called by the CLI, the web UI, and can be called by the Orchestrator
to trigger a refresh after new documents are added.
Args:
docs_dir: Path to the folder containing your documents.
force: Re-index everything even if files haven't changed.
"""
indexer = _get_indexer()
indexer.index_directory(docs_dir=docs_dir, force=force)
def index_single(path: str):
"""Index a single file immediately (called by file-watcher or API upload)."""
indexer = _get_indexer()
indexer.index_single_file(path)
def get_knowledge_stats() -> Dict:
"""Return stats about the current knowledge base (for health checks)."""
return _get_indexer().get_stats()
def list_indexed_docs() -> List[Dict]:
"""List all documents currently indexed."""
return _get_indexer().list_documents()
def delete_doc(filename: str):
"""Remove a document and all its chunks from the index."""
_get_indexer().delete_document(filename)
def stream_knowledge_answer(question: str, top_k: int = 5):
"""
Streaming variant — yields tokens one by one.
Used by the FastAPI /stream endpoint for the web UI.
"""
indexer = _get_indexer()
chunks = indexer.query(question, top_k=top_k)
if not chunks:
yield "No documents indexed yet."
return
yield from stream_answer(question, chunks) |