Spaces:
Sleeping
Sleeping
File size: 3,585 Bytes
2d25973 | 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 | import os
import uuid
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from google import genai
from google.genai import types
import streamlit as st
# Initialize ChromaDB persistent client locally
CHROMA_DB_DIR = "./chroma_db"
os.makedirs(CHROMA_DB_DIR, exist_ok=True)
@st.cache_resource
def get_chroma_client():
return chromadb.PersistentClient(path=CHROMA_DB_DIR)
def get_gemini_embedding(text):
"""Hits the Gemini API to mathematically embed a text chunk into vectors."""
client = genai.Client()
response = client.models.embed_content(
model='gemini-embedding-001',
contents=text,
)
return response.embeddings[0].values
def chunk_text(text):
"""Splits massive PDFs into smaller, 1000-character overlapping chunks for precision RAG."""
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
length_function=len,
is_separator_regex=False,
)
return text_splitter.split_text(text)
def embed_and_store_document(user_id, document_id, pdf_text, filename):
"""Chunks a raw PDF, calculates vectors for each chunk, and saves to the user's Chroma collection."""
client = get_chroma_client()
collection = client.get_or_create_collection(name="user_libraries")
# Check if this document was already embedded to skip duplicate processing
existing = collection.get(where={"document_id": document_id})
if existing and existing['ids']:
return # Already indexed
chunks = chunk_text(pdf_text)
ids = []
embeddings = []
documents = []
metadatas = []
for i, chunk in enumerate(chunks):
if not chunk.strip():
continue
try:
emb = get_gemini_embedding(chunk)
ids.append(f"{document_id}_chunk_{i}")
embeddings.append(emb)
documents.append(chunk)
metadatas.append({"user_id": user_id, "document_id": document_id, "filename": filename})
except Exception as e:
print(f"Failed to embed chunk {i}: {e}")
if ids:
collection.add(
embeddings=embeddings,
documents=documents,
metadatas=metadatas,
ids=ids
)
def query_relevant_chunks(user_id, query, n_results=5):
"""Full Library Search: Searches ALL textbooks uploaded by the user to find the best 5 paragraphs."""
client = get_chroma_client()
collection = client.get_or_create_collection(name="user_libraries")
if collection.count() == 0:
return ""
try:
query_embedding = get_gemini_embedding(query)
results = collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
where={"user_id": user_id} # Filter exclusively to this user's entire library!
)
if not results['documents'] or not results['documents'][0]:
return ""
# Compile all matching text chunks into a unified context string, tagging the source book
context_pieces = []
for i, text in enumerate(results['documents'][0]):
source_file = results['metadatas'][0][i].get("filename", "Unknown Document")
context_pieces.append(f"[Excerpt from {source_file}]:\n{text}")
return "\n\n---\n\n".join(context_pieces)
except Exception as e:
print(f"Semantic search failed: {e}")
return ""
|