File size: 2,617 Bytes
27551c8 | 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 | import chromadb
from chromadb.utils import embedding_functions
import os
class BaitakMemory:
def __init__(self, collection_name="baitak_insights", path="/home/ubuntu/Project-Baitak-Intelligence/data/chroma_db"):
self.client = chromadb.PersistentClient(path=path)
# Using OpenAI embeddings for simplicity, can be replaced with Gemini if API key is available
self.collection = self.client.get_or_create_collection(name=collection_name)
def add_insight(self, insight_text: str, metadata: dict = None):
# ChromaDB requires a unique ID for each document
doc_id = f"insight_{self.collection.count() + 1}"
self.collection.add(
documents=[insight_text],
metadatas=[metadata if metadata else {}],
ids=[doc_id]
)
print(f"Added insight with ID: {doc_id}")
def query_insights(self, query_text: str, n_results: int = 5):
results = self.collection.query(
query_texts=[query_text],
n_results=n_results
)
return results
def get_all_insights(self):
return self.collection.get(include=['documents', 'metadatas'])
def clear_memory(self):
self.client.delete_collection(name=self.collection.name)
self.collection = self.client.get_or_create_collection(name=self.collection.name, embedding_function=self.embedding_function)
print(f"Memory collection \'{self.collection.name}\' cleared.")
if __name__ == "__main__":
memory = BaitakMemory()
memory.clear_memory() # Clear for fresh start
# Add some synthetic insights
memory.add_insight("KFH's Net Financing Margin (NFM) remained stable around 3.5% in FY2025 due to CBK rate cuts.", {"quarter": "Q4-2025", "kpi": "NFM"})
memory.add_insight("Digital transformation initiatives led to a 10% increase in mobile banking users in Q3-2025.", {"quarter": "Q3-2025", "kpi": "Digital Transformation"})
memory.add_insight("The Cost-to-Income ratio improved to 34.06% in FY2025, down from 35.46% in FY2024.", {"quarter": "FY2025", "kpi": "Cost-to-Income"})
# Query insights
print("\nQuerying for NFM insights:")
results = memory.query_insights("What was the Net Financing Margin in 2025?")
for doc, meta in zip(results["documents"], results["metadatas"]):
print(f"Insight: {doc}, Metadata: {meta}")
print("\nQuerying for digital transformation:")
results = memory.query_insights("Tell me about digital milestones.")
for doc, meta in zip(results["documents"], results["metadatas"]):
print(f"Insight: {doc}, Metadata: {meta}")
|