Narrow the size of the dataframe returned by DatasetSemanticSearchAgent

agents/dataset_semantic_search_agent.py

@@ -7,7 +7,7 @@ class DatasetSemanticSearchAgent:
     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:
+    def run(self, query: str, limit: int = 10, score_threshold: int = 1.2) -> tuple[pd.DataFrame, dict[str, float]]:
         docs_with_scores = self._vector_store.similarity_search_with_score(query, k=limit)
 
         results = []
@@ -24,4 +24,6 @@ class DatasetSemanticSearchAgent:
         df = pd.DataFrame(results)
         df.rename(columns={"_id": "id", "full_text": "content"}, inplace=True)
 
-        return df
+        df = df[["title", "author", "date", "url", "content"]]
+
+        return df, {"cost": 0}