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:
parent
5b3f4bd8c0
commit
315cd77639
@ -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"
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|||||||
@ -16,6 +16,8 @@ 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_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_list_collections, handler = rag.list_collections ),
|
||||||
gadget.tools_mapping( schema = rag.schema_inspect_collection, handler = rag.inspect_collection ),
|
gadget.tools_mapping( schema = rag.schema_inspect_collection, handler = rag.inspect_collection ),
|
||||||
]
|
]
|
||||||
|
|||||||
@ -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)
|
||||||
|
|
||||||
|
|||||||
118
tools/rag.py
118
tools/rag.py
@ -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}"
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user