feat: hybrid collection management — agent can create/delete/list collections via prompt, remove RAG_COLLECTIONS config, switch to ONNX MiniLM embedding (local, no API)

This commit is contained in:
Dita Aji Pratama 2026-05-25 10:25:06 +07:00
parent 5b3f4bd8c0
commit 315cd77639
4 changed files with 84 additions and 56 deletions

View File

@ -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" ) ) MAX_TOOL_OUTPUT = int( os.getenv("MAX_TOOL_OUTPUT", default="4000" ) )
# RAG Configuration # RAG Configuration
RAG_PERSIST_DIR = os.getenv("RAG_PERSIST_DIR", default="chroma_db" ) RAG_PERSIST_DIR = os.getenv("RAG_PERSIST_DIR", default="chroma_db" )
RAG_EMBEDDING_MODEL = os.getenv("RAG_EMBEDDING_MODEL", default="nomic-embed-text" ) # Embedding: ChromaDB ONNX default (all-MiniLM-L6-v2, lokal, tidak perlu API call)
RAG_COLLECTIONS = {
"food_recommendations": {
"description": "Menu makanan, preferensi pelanggan, data kuliner"
},
}

View File

@ -16,8 +16,10 @@ tools_definition = [
gadget.tools_mapping( schema = coder.schema_git_operation, handler = coder.git_operation ), 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_store_knowledge, handler = rag.store_knowledge ),
gadget.tools_mapping( schema = rag.schema_search_knowledge, handler = rag.search_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_create_collection, handler = rag.create_collection ),
gadget.tools_mapping( schema = rag.schema_inspect_collection, handler = rag.inspect_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 # Ekstrak dari tools_definition ke dua format berbeda

View File

@ -34,13 +34,16 @@ def build_system_prompt(tools_definition):
"All file operations are relative to this directory.", "All file operations are relative to this directory.",
"", "",
"RAG capabilities (knowledge retrieval):", "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.", "- inspect_collection → learn metadata fields before searching.",
"- search_knowledge → semantic search + optional metadata filter.", "- search_knowledge → semantic search + optional metadata filter.",
"- store_knowledge → save docs with rich metadata for later use.", "- store_knowledge → save docs with rich metadata for later use.",
"", "",
"RAG workflow: inspect → search → reason. Always inspect a collection", "You can create collections yourself! When you encounter a new topic,",
"first to discover its metadata keys, then use them in search filters." "use create_collection first, then store_knowledge to populate it.",
"Always inspect_collection to discover metadata keys before filtering."
]) ])
return "\n".join(lines) return "\n".join(lines)

View File

@ -1,42 +1,13 @@
import json import json
import urllib.request
import urllib.error
from urllib.parse import urlparse
import chromadb import chromadb
from chromadb.config import Settings from chromadb.config import Settings
import config 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 ─────────────────────────────────────────────── # ── ChromaDB singleton ───────────────────────────────────────────────
_store = None _store = None
_ef = None
def _get_store(): def _get_store():
global _store global _store
@ -47,17 +18,9 @@ def _get_store():
) )
return _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): def _collection(name):
if name not in config.RAG_COLLECTIONS: """Get or create collection — uses ChromaDB's default ONNX embedding (all-MiniLM-L6-v2)."""
avail = ", ".join(config.RAG_COLLECTIONS) return _get_store().get_or_create_collection(name=name)
raise ValueError(f"Unknown collection '{name}'. Available: {avail}")
return _get_store().get_or_create_collection(name=name, embedding_function=_get_ef())
# ── Tool schemas ───────────────────────────────────────────────────── # ── 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 = { schema_list_collections = {
"type": "function", "type": "function",
"function": { "function": {
"name": "list_collections", "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": {}} "parameters": {"type": "object", "properties": {}}
} }
} }
@ -231,13 +238,34 @@ def search_knowledge(collection, query, n_results=5, filter=None):
return f"Error: {e}" 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(): def list_collections():
try: try:
if not config.RAG_COLLECTIONS: cols = _get_store().list_collections()
return "No collections defined in config." if not cols:
return "Available collections:\n" + "\n".join( return "No collections exist yet."
f"- {n}: {i.get('description', '')}" for n, i in config.RAG_COLLECTIONS.items() 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: except Exception as e:
return f"Error: {e}" return f"Error: {e}"