cynthesis-v4 / semantic_search_agent.py
abdoh-alkhateeb's picture
Update method name of SemanticSearchAgent for consistency
f08a437
raw
history blame
No virus
946 Bytes
import pandas as pd
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
class SemanticSearchAgent:
def __init__(self, vector_store_path: str) -> None:
self._vector_store = FAISS.load_local(vector_store_path, HuggingFaceEmbeddings(), allow_dangerous_deserialization=True)
def run(self, query: str, limit: int = 10, score_threshold: int = 1.2) -> pd.DataFrame:
docs_with_scores = self._vector_store.similarity_search_with_score(query, k=limit)
results = []
for doc, score in docs_with_scores:
if score >= score_threshold:
continue
result = doc.metadata
result["summary"] = doc.page_content
result["score"] = score
results.append(result)
df = pd.DataFrame(results)
df.rename(columns={"_id": "id", "full_text": "content"}, inplace=True)
return df