import json import urllib.request import urllib.error from urllib.parse import urlparse import chromadb from chromadb.config import Settings import config # ── Embedding (Ollama) ─────────────────────────────────────────────── from chromadb.api.types import EmbeddingFunction, Embeddings class OllamaEmbeddingFunction(EmbeddingFunction): def __init__(self, base_url, model): parsed = urlparse(base_url.rstrip('/')) self.ollama_base = f"{parsed.scheme}://{parsed.netloc}" self.model = model def __call__(self, input) -> Embeddings: url = f"{self.ollama_base}/api/embed" texts = input if isinstance(input, list) else [input] payload = {"model": self.model, "input": texts} data = json.dumps(payload).encode('utf-8') req = urllib.request.Request(url, data=data, method='POST') req.add_header('Content-Type', 'application/json') try: with urllib.request.urlopen(req, timeout=30) as resp: response = json.loads(resp.read().decode('utf-8')) return response["embeddings"] except Exception as e: raise RuntimeError(f"Embedding error: {e}") # ── ChromaDB singleton ─────────────────────────────────────────────── _store = None _ef = None def _get_store(): global _store if _store is None: _store = chromadb.PersistentClient( path=config.RAG_PERSIST_DIR, settings=Settings(anonymized_telemetry=False), ) return _store def _get_ef(): global _ef if _ef is None: _ef = OllamaEmbeddingFunction(config.llm_baseurl, config.RAG_EMBEDDING_MODEL) return _ef def _collection(name): if name not in config.RAG_COLLECTIONS: avail = ", ".join(config.RAG_COLLECTIONS) raise ValueError(f"Unknown collection '{name}'. Available: {avail}") return _get_store().get_or_create_collection(name=name, embedding_function=_get_ef()) # ── Tool schemas ───────────────────────────────────────────────────── schema_store_knowledge = { "type": "function", "function": { "name": "store_knowledge", "description": ( "Store one or more documents with arbitrary metadata into a RAG collection. " "Metadata is a free-form dict — choose meaningful keys for future filtering " "(e.g., restaurant, category, allergens, spice_level, taste_profile, price" ", customer_id, dietary)." ), "parameters": { "type": "object", "properties": { "collection": { "type": "string", "description": "Target collection name (must be defined in config)" }, "documents": { "type": "array", "items": { "type": "object", "properties": { "id": {"type": "string", "description": "Unique document ID"}, "text": {"type": "string", "description": "Document body text"}, "metadata": { "type": "object", "description": "Arbitrary key-value metadata", "default": {} } }, "required": ["id", "text"] }, "description": "List of documents to persist" } }, "required": ["collection", "documents"] } } } schema_search_knowledge = { "type": "function", "function": { "name": "search_knowledge", "description": ( "Semantically search a RAG collection. Optionally narrow with a " "metadata filter using ChromaDB where syntax. " "Examples: {'category': 'main_course'}, {'spice_level': {'$lte': 2}}, " "{'allergens': {'$contains': 'seafood'}}." ), "parameters": { "type": "object", "properties": { "collection": { "type": "string", "description": "Collection name to search in" }, "query": { "type": "string", "description": "Natural-language search query" }, "n_results": { "type": "integer", "description": "Max results to return (default 5)", "default": 5 }, "filter": { "type": "object", "description": "Optional metadata filter dict", "default": None } }, "required": ["collection", "query"] } } } schema_list_collections = { "type": "function", "function": { "name": "list_collections", "description": "List all available RAG collections defined in config with their descriptions.", "parameters": {"type": "object", "properties": {}} } } schema_inspect_collection = { "type": "function", "function": { "name": "inspect_collection", "description": ( "Examine sample documents and metadata fields in a RAG collection. " "Always call this before search_knowledge to learn what metadata keys " "are available for filtering, then pass them in the filter parameter." ), "parameters": { "type": "object", "properties": { "collection": { "type": "string", "description": "Collection name to inspect" }, "sample_size": { "type": "integer", "description": "Number of sample documents (default 3)", "default": 3 } }, "required": ["collection"] } } } # ── Tool handlers ──────────────────────────────────────────────────── def _sanitize_meta(meta): """ChromaDB metadata only allows str/int/float/bool. Convert lists to JSON string, remove empty lists.""" out = {} for k, v in meta.items(): if isinstance(v, list): if len(v) == 0: continue out[k] = json.dumps(v, ensure_ascii=False) elif isinstance(v, (str, int, float, bool)): out[k] = v else: out[k] = str(v) return out def store_knowledge(collection, documents): try: col = _collection(collection) ids, texts, metas = [], [], [] for doc in documents: ids.append(doc["id"]) texts.append(doc["text"]) metas.append(_sanitize_meta(doc.get("metadata", {}))) col.add(ids=ids, documents=texts, metadatas=metas) return f"Stored {len(documents)} document(s) in '{collection}'." except Exception as e: return f"Error: {e}" def search_knowledge(collection, query, n_results=5, filter=None): try: col = _collection(collection) kw = {"query_texts": [query], "n_results": n_results} if filter: kw["where"] = filter r = col.query(**kw) if not r["ids"] or not r["ids"][0]: return "No results found." out = [] for i in range(len(r["ids"][0])): did = r["ids"][0][i] txt = r["documents"][0][i] if len(txt) > 500: txt = txt[:500] + "..." meta = json.dumps(r["metadatas"][0][i], ensure_ascii=False) if r.get("metadatas") else "{}" dist = "" if r.get("distances"): dist = f" (score: {r['distances'][0][i]:.4f})" out.append(f"[{did}]{dist}\n text: {txt}\n metadata: {meta}") return "\n---\n".join(out) except Exception as e: return f"Error: {e}" def list_collections(): try: if not config.RAG_COLLECTIONS: return "No collections defined in config." return "Available collections:\n" + "\n".join( f"- {n}: {i.get('description', '')}" for n, i in config.RAG_COLLECTIONS.items() ) except Exception as e: return f"Error: {e}" def inspect_collection(collection, sample_size=3): try: col = _collection(collection) cnt = col.count() if cnt == 0: return f"Collection '{collection}' is empty." n = min(sample_size, cnt) r = col.get(limit=n, include=["documents", "metadatas"]) out = [f"Collection: {collection} | Total documents: {cnt}", f"Sample ({n}):"] for i in range(len(r["ids"])): txt = r["documents"][i] if len(txt) > 200: txt = txt[:200] + "..." meta = json.dumps(r["metadatas"][i], ensure_ascii=False) if r.get("metadatas") and r["metadatas"][i] else "(none)" out.append(f"\n [{r['ids'][i]}] text: {txt} metadata: {meta}") keys = set() for m in r["metadatas"]: if m: keys.update(m.keys()) if keys: out.append(f"\nMetadata keys: {', '.join(sorted(keys))}") return "\n".join(out) except Exception as e: return f"Error: {e}"