File size: 5,456 Bytes
d3ab573 | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | import chromadb
from sklearn.cluster import KMeans
import pandas as pd
import json
import ast
from itertools import groupby
import uuid
from datetime import datetime
from pathlib import Path
root_dir = Path(__file__).parent.parent
data_dir = root_dir / "datasets"
embed_faq_df = pd.read_csv(data_dir / "processed" / "bank_faq" / "embed_faq.csv")
# chroma db
client = chromadb.PersistentClient(path="./chroma_db")
# client.delete_collection("bank_faq")
# client.delete_collection("chat_history")
# create collection
collection = client.get_or_create_collection(
name="bank_faq",
metadata={"hnsw:space": "cosine"} # cosine similarity for sentence transformers
)
embed_faq_df["embedding"] = embed_faq_df["embedding"].apply(ast.literal_eval)
# prepare data
# ids = [str(i) for i in range(len(embed_faq_df))]
ids = [
f"doc_{row['source_id']}_chunk_{row['chunk_num']}"
for _, row in embed_faq_df.iterrows()
]
embeddings = embed_faq_df["embedding"].tolist()
documents = embed_faq_df["text"].tolist()
metadatas = [
{
"doc_id": f"doc_{row['source_id']}_chunk_{row['chunk_num']}",
"domain": row["domain"],
"source_id": row["source_id"],
"chunk_num": row["chunk_num"],
"question": row["question"]
}
for _, row in embed_faq_df.iterrows()
]
# print(ids)
# print(type(embeddings[0]))
# print(documents)
# print(metadatas)
# store
collection.add(
ids=ids,
embeddings=embeddings,
documents=documents,
metadatas=metadatas
)
print(f"Stored {collection.count()} documents")
faq_collection = client.get_collection(name="bank_faq")
results = faq_collection.get(
include=["documents", "metadatas"]
)
source_id_to_doc_ids = {}
for doc_id, meta in zip(ids, metadatas):
source_id = meta["source_id"]
if source_id not in source_id_to_doc_ids:
source_id_to_doc_ids[source_id] = []
source_id_to_doc_ids[source_id].append(doc_id)
# build a dict grouped by source_id
groups = {}
for doc_id, meta, text in zip(results["ids"], results["metadatas"], results["documents"]):
source_id = meta["source_id"]
if source_id not in groups:
groups[source_id] = []
groups[source_id].append({
"doc_id": doc_id,
"chunk_num": int(meta["chunk_num"]),
"domain": meta["domain"],
"question": meta["question"],
"text": text.split("A:", 1)[1].strip() if "A:" in text else text,
})
# build viewable
viewable = []
for source_id, chunks in groups.items():
chunks = sorted(chunks, key=lambda x: x["chunk_num"]) # sort by chunk order
entry = {
"id": source_id,
"domain": chunks[0]["domain"],
"question": chunks[0]["question"],
}
if len(chunks) == 1:
entry["answer"] = chunks[0]["text"]
else:
entry["answer"] = {"chunks": [{"chunk": c["chunk_num"] + 1, "doc_id": c["doc_id"], "text": c["text"]} for c in chunks]}
viewable.append(entry)
with open("chroma_preview.json", "w") as f:
json.dump(viewable, f, indent=2)
print("Saved to chroma_preview.json")
with open("chroma_preview.json", "r") as f: # or test.json / val.json
faq_db = json.load(f)
for item in faq_db:
source_id = item["id"]
item["relevant_doc_ids"] = source_id_to_doc_ids.get(source_id, [])
with open("chroma_preview.json", "w") as f:
json.dump(faq_db, f, indent=2)
# save chats in db
chat_collection = client.get_or_create_collection(name="chat_history")
def create_new_chat():
# creates a new chat session and returns its ID
chat_id = str(uuid.uuid4())
print(f"New chat created: {chat_id}")
return chat_id
def get_chat_history(chat_id):
"""Retrieve full history of a specific chat."""
results = chat_collection.get(
where={"chat_id": chat_id}
)
# Pair up queries and answers
history = []
for doc, meta in zip(results["documents"], results["metadatas"]):
history.append({
"query": doc,
"answer": meta["answer"],
"timestamp": meta["timestamp"]
})
# Sort by timestamp
history.sort(key=lambda x: x["timestamp"])
return history
def load_chats_from_db():
# rebuild the chats dict from ChromaDB on page load
try:
results = chat_collection.get(include=["metadatas", "documents"])
except Exception:
return {}
chats = {}
for doc, meta in zip(results["documents"], results["metadatas"]):
chat_id = meta.get("chat_id")
chat_name = meta.get("chat_name", "Chat 1")
timestamp = meta.get("timestamp", "")
query = doc
answer = meta.get("answer", "")
if chat_name not in chats:
chats[chat_name] = {"messages": [], "chat_id": chat_id}
chats[chat_name]["messages"].append({"role": "user", "content": query, "timestamp": timestamp})
chats[chat_name]["messages"].append({"role": "assistant", "content": answer, "timestamp": timestamp})
# sort messages within each chat by timestamp
for chat_name in chats:
chats[chat_name]["messages"].sort(key=lambda x: x.get("timestamp", ""))
return chats
def delete_chat(chat_id: str):
"""Remove a chat from ChromaDB and the JSON log."""
# Remove from ChromaDB
results = chat_collection.get(where={"chat_id": chat_id})
if results["ids"]:
chat_collection.delete(ids=results["ids"])
print(f"Chat deleted: {chat_id}")
|