From 315cd776399773441a511076993d56b0b4155782 Mon Sep 17 00:00:00 2001 From: Dita Aji Pratama Date: Mon, 25 May 2026 10:25:06 +0700 Subject: [PATCH] =?UTF-8?q?feat:=20hybrid=20collection=20management=20?= =?UTF-8?q?=E2=80=94=20agent=20can=20create/delete/list=20collections=20vi?= =?UTF-8?q?a=20prompt,=20remove=20RAG=5FCOLLECTIONS=20config,=20switch=20t?= =?UTF-8?q?o=20ONNX=20MiniLM=20embedding=20(local,=20no=20API)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config.py | 7 +-- hendrik.py | 6 ++- scripts/gadget.py | 9 ++-- tools/rag.py | 118 ++++++++++++++++++++++++++++------------------ 4 files changed, 84 insertions(+), 56 deletions(-) diff --git a/config.py b/config.py index f656478..930667d 100644 --- a/config.py +++ b/config.py @@ -14,9 +14,4 @@ AGENT_MAX_ITERATIONS = int( os.getenv("AGENT_MAX_ITERATIONS", default="10" MAX_TOOL_OUTPUT = int( os.getenv("MAX_TOOL_OUTPUT", default="4000" ) ) # RAG Configuration RAG_PERSIST_DIR = os.getenv("RAG_PERSIST_DIR", default="chroma_db" ) -RAG_EMBEDDING_MODEL = os.getenv("RAG_EMBEDDING_MODEL", default="nomic-embed-text" ) -RAG_COLLECTIONS = { - "food_recommendations": { - "description": "Menu makanan, preferensi pelanggan, data kuliner" - }, -} +# Embedding: ChromaDB ONNX default (all-MiniLM-L6-v2, lokal, tidak perlu API call) diff --git a/hendrik.py b/hendrik.py index 0490211..e6be11c 100644 --- a/hendrik.py +++ b/hendrik.py @@ -16,8 +16,10 @@ tools_definition = [ gadget.tools_mapping( schema = coder.schema_git_operation, handler = coder.git_operation ), gadget.tools_mapping( schema = rag.schema_store_knowledge, handler = rag.store_knowledge ), gadget.tools_mapping( schema = rag.schema_search_knowledge, handler = rag.search_knowledge ), - gadget.tools_mapping( schema = rag.schema_list_collections, handler = rag.list_collections ), - gadget.tools_mapping( schema = rag.schema_inspect_collection, handler = rag.inspect_collection ), + gadget.tools_mapping( schema = rag.schema_create_collection, handler = rag.create_collection ), + gadget.tools_mapping( schema = rag.schema_delete_collection, handler = rag.delete_collection ), + gadget.tools_mapping( schema = rag.schema_list_collections, handler = rag.list_collections ), + gadget.tools_mapping( schema = rag.schema_inspect_collection, handler = rag.inspect_collection ), ] # Ekstrak dari tools_definition ke dua format berbeda diff --git a/scripts/gadget.py b/scripts/gadget.py index e8a5d6a..d9f1031 100644 --- a/scripts/gadget.py +++ b/scripts/gadget.py @@ -34,13 +34,16 @@ def build_system_prompt(tools_definition): "All file operations are relative to this directory.", "", "RAG capabilities (knowledge retrieval):", - "- list_collections → see available knowledge bases.", + "- list_collections → see available collections & doc counts.", + "- create_collection → create a new collection for a new topic.", + "- delete_collection → permanently remove a collection and its data.", "- inspect_collection → learn metadata fields before searching.", "- search_knowledge → semantic search + optional metadata filter.", "- store_knowledge → save docs with rich metadata for later use.", "", - "RAG workflow: inspect → search → reason. Always inspect a collection", - "first to discover its metadata keys, then use them in search filters." + "You can create collections yourself! When you encounter a new topic,", + "use create_collection first, then store_knowledge to populate it.", + "Always inspect_collection to discover metadata keys before filtering." ]) return "\n".join(lines) diff --git a/tools/rag.py b/tools/rag.py index 312db55..2438143 100644 --- a/tools/rag.py +++ b/tools/rag.py @@ -1,42 +1,13 @@ 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 @@ -47,17 +18,9 @@ def _get_store(): ) 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()) + """Get or create collection — uses ChromaDB's default ONNX embedding (all-MiniLM-L6-v2).""" + return _get_store().get_or_create_collection(name=name) # ── Tool schemas ───────────────────────────────────────────────────── @@ -139,11 +102,55 @@ schema_search_knowledge = { } } +schema_create_collection = { + "type": "function", + "function": { + "name": "create_collection", + "description": ( + "Create a new RAG collection for a new topic/domain. Use a short, descriptive name " + "with underscores (e.g., 'tanaman_hias', 'customer_profiles'). Optionally provide a description." + ), + "parameters": { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Collection name (lowercase, underscores for spaces)" + }, + "description": { + "type": "string", + "description": "What this collection stores", + "default": "" + } + }, + "required": ["name"] + } + } +} + +schema_delete_collection = { + "type": "function", + "function": { + "name": "delete_collection", + "description": "Permanently delete an entire RAG collection and all documents in it. Use with caution — this cannot be undone.", + "parameters": { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Collection name to delete" + } + }, + "required": ["name"] + } + } +} + schema_list_collections = { "type": "function", "function": { "name": "list_collections", - "description": "List all available RAG collections defined in config with their descriptions.", + "description": "List all existing RAG collections with their document count and description.", "parameters": {"type": "object", "properties": {}} } } @@ -231,13 +238,34 @@ def search_knowledge(collection, query, n_results=5, filter=None): return f"Error: {e}" +def create_collection(name, description=""): + try: + col = _get_store().get_or_create_collection(name=name) + col.modify(metadata={"description": description}) + return f"Collection '{name}' is ready." + except Exception as e: + return f"Error: {e}" + +def delete_collection(name): + try: + _get_store().delete_collection(name) + return f"Deleted collection '{name}'." + 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() - ) + cols = _get_store().list_collections() + if not cols: + return "No collections exist yet." + out = ["Available collections:"] + for col in cols: + meta = col.metadata or {} + desc = meta.get("description", "") + cnt = col.count() + tag = f" ({desc})" if desc else "" + out.append(f"- {col.name}{tag} [{cnt} docs]") + return "\n".join(out) except Exception as e: return f"Error: {e}"