Spaces:
Running
Running
File size: 6,370 Bytes
9a1014e | 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 | from __future__ import annotations
import os
import numpy as np
from dotenv import load_dotenv
from langchain_community.vectorstores import Chroma
from hf_text_embeddings import HFTextEmbeddings
DOCUMENT_PERSIST_DIR = "./chroma_db_docs"
DOCUMENT_COLLECTION = "document_context"
load_dotenv()
def _metadata_as_text(metadata: dict) -> str:
searchable_fields = (
"document_title",
"document_type",
"document_class",
"keywords",
"filename",
)
return " | ".join(
str(metadata[field]).strip()
for field in searchable_fields
if metadata.get(field)
)
def _combine_filters(first: dict | None, second: dict) -> dict:
if not first:
return second
return {"$and": [first, second]}
def _classify_document_filter(
vectorstore: Chroma,
embeddings: HFTextEmbeddings,
query: str,
base_filter: dict | None = None,
) -> tuple[dict | None, dict | None, float | None]:
get_kwargs = {"include": ["metadatas"]}
if base_filter:
get_kwargs["where"] = base_filter
stored = vectorstore.get(**get_kwargs)
candidates: dict[str, dict] = {}
for metadata in stored.get("metadatas") or []:
if not metadata:
continue
source = str(metadata.get("source") or "")
if source and _metadata_as_text(metadata):
candidates.setdefault(source, metadata)
if not candidates:
return base_filter, None, None
sources = list(candidates)
metadata_texts = [_metadata_as_text(candidates[source]) for source in sources]
query_vector = np.asarray(embeddings.embed_query(query), dtype=np.float32)
metadata_vectors = np.asarray(
embeddings.embed_documents(metadata_texts),
dtype=np.float32,
)
query_norm = np.linalg.norm(query_vector)
metadata_norms = np.linalg.norm(metadata_vectors, axis=1)
denominators = metadata_norms * query_norm
scores = np.divide(
metadata_vectors @ query_vector,
denominators,
out=np.zeros(len(metadata_vectors), dtype=np.float32),
where=denominators != 0,
)
best_index = int(np.argmax(scores))
best_source = sources[best_index]
best_metadata = candidates[best_source]
return (
_combine_filters(base_filter, {"source": {"$eq": best_source}}),
best_metadata,
float(scores[best_index]),
)
def retrieve_document_context(
query: str,
k: int = 4,
filter_metadata: dict | None = None,
) -> str:
try:
embeddings = HFTextEmbeddings()
vectorstore = Chroma(
collection_name=os.getenv("DOCUMENT_CHROMA_COLLECTION", DOCUMENT_COLLECTION),
embedding_function=embeddings,
persist_directory=os.getenv("DOCUMENT_CHROMA_DIR", DOCUMENT_PERSIST_DIR),
)
selected_filter, selected_metadata, metadata_score = _classify_document_filter(
vectorstore,
embeddings,
query,
filter_metadata,
)
if selected_metadata is not None:
print(
"Metadata classifier | "
f"score={metadata_score:.4f} | "
f"document={selected_metadata.get('document_title')} | "
f"type={selected_metadata.get('document_type')}",
flush=True,
)
# Evita el query HNSW de Chroma, que puede bloquearse en algunos
# contenedores. La coleccion ASTM es pequena, por lo que un ranking
# coseno directo sobre los embeddings persistidos es rapido y estable.
get_kwargs = {"include": ["documents", "metadatas", "embeddings"]}
if selected_filter:
get_kwargs["where"] = selected_filter
stored = vectorstore.get(**get_kwargs)
documents = stored.get("documents") or []
metadatas = stored.get("metadatas") or []
stored_embeddings = stored.get("embeddings")
if not documents or stored_embeddings is None:
print("Document Chroma returned no searchable chunks.", flush=True)
return ""
query_vector = np.asarray(embeddings.embed_query(query), dtype=np.float32)
chunk_vectors = np.asarray(stored_embeddings, dtype=np.float32)
query_norm = np.linalg.norm(query_vector)
chunk_norms = np.linalg.norm(chunk_vectors, axis=1)
denominators = chunk_norms * query_norm
similarities = np.divide(
chunk_vectors @ query_vector,
denominators,
out=np.zeros(len(chunk_vectors), dtype=np.float32),
where=denominators != 0,
)
top_indices = np.argsort(similarities)[::-1][: min(k, len(documents))]
blocks = []
for rank, index in enumerate(top_indices, start=1):
metadata = metadatas[index] if index < len(metadatas) else {}
print(
f"Document rank {rank} | similarity={similarities[index]:.4f} | "
f"document={metadata.get('document_title')}",
flush=True,
)
blocks.append(
f"[Documento {rank}] metadata={metadata}\n{documents[index]}"
)
return "\n\n".join(blocks)
except Exception as exc:
print(f"Document Chroma retrieval skipped: {exc}", flush=True)
return ""
def build_context_prompt(
user_text: str,
k: int = 4,
conversation_history: str = "",
) -> str:
document_context = retrieve_document_context(user_text, k=k)
context_parts = []
if document_context:
context_parts.append("Contexto normativo/documental recuperado de documentos/Chroma:\n" + document_context)
if not context_parts:
return user_text
context = "\n\n".join(context_parts)
return f"""Sos el agente del proyecto Alberti/Metalurgia.
Usa como fuente principal el contexto técnico y normativo recuperado desde los documentos indexados en Chroma.
Cuando respondas sobre datos del proyecto, prioriza los registros recuperados y menciona información relacionada si ayudan pero no divulgues ID o datos propios de los registros.
Si el contexto recuperado no alcanza para responder con precision, dilo claramente y no inventes datos exactos.
{context}
Historial reciente de esta conversacion:
{conversation_history or "Sin mensajes anteriores."}
Pregunta del usuario:
{user_text}"""
|