Spaces:
Sleeping
Sleeping
File size: 5,674 Bytes
b27eb78 |
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 |
import json
from pathlib import Path
import yaml
from loguru import logger
from opik import opik_context, track
from smolagents import Tool
from second_brain_online.application.rag import get_retriever
class MongoDBRetrieverTool(Tool):
name = "mongodb_vector_search_retriever"
description = """Use this tool to search and retrieve relevant documents from a knowledge base using semantic search.
This tool performs similarity-based search to find the most relevant documents matching the query.
Best used when you need to:
- Find specific information from stored documents
- Get context about a topic
- Research historical data or documentation
The tool will return multiple relevant document snippets."""
inputs = {
"query": {
"type": "string",
"description": """The search query to find relevant documents for using semantic search.
Should be a clear, specific question or statement about the information you're looking for.""",
}
}
output_type = "string"
def __init__(self, config_path: Path, **kwargs):
super().__init__(**kwargs)
self.config_path = config_path
self.retriever = self.__load_retriever(config_path)
def __load_retriever(self, config_path: Path):
config = yaml.safe_load(config_path.read_text())
config = config["parameters"]
return get_retriever(
embedding_model_id=config["embedding_model_id"],
embedding_model_type=config["embedding_model_type"],
retriever_type=config["retriever_type"],
k=5,
device=config["device"],
)
@track(name="MongoDBRetrieverTool.forward")
def forward(self, query: str) -> str:
if hasattr(self.retriever, "search_kwargs"):
search_kwargs = self.retriever.search_kwargs
else:
try:
search_kwargs = {
"fulltext_penalty": self.retriever.fulltext_penalty,
"vector_score_penalty": self.retriever.vector_penalty,
"top_k": self.retriever.top_k,
}
except AttributeError:
logger.warning("Could not extract search kwargs from retriever.")
search_kwargs = {}
opik_context.update_current_trace(
tags=["agent"],
metadata={
"search": search_kwargs,
"embedding_model_id": self.retriever.vectorstore.embeddings.model,
},
)
try:
query = self.__parse_query(query)
relevant_docs = self.retriever.invoke(query)
formatted_docs = []
for i, doc in enumerate(relevant_docs, 1):
# Extract metadata
title = doc.metadata.get("title", "Untitled")
datetime = doc.metadata.get("datetime", "unknown")
contextual_summary = doc.metadata.get("contextual_summary", "")
marketing_insights = doc.metadata.get("marketing_insights", {})
content = doc.page_content.strip()
# Format marketing insights if available
marketing_insights_text = ""
if marketing_insights:
marketing_insights_text = "\n<marketing_insights>\n"
# Add quotes
quotes = marketing_insights.get("quotes", [])
if quotes:
marketing_insights_text += "<quotes>\n"
for quote in quotes:
marketing_insights_text += f"- \"{quote.get('quote', '')}\" (Sentiment: {quote.get('sentiment', 'Unknown')})\n"
marketing_insights_text += "</quotes>\n"
# Add key findings
findings = marketing_insights.get("key_findings", [])
if findings:
marketing_insights_text += "<key_findings>\n"
for finding in findings:
marketing_insights_text += f"- {finding.get('finding', '')} (Impact: {finding.get('impact', 'Unknown')})\n"
marketing_insights_text += "</key_findings>\n"
marketing_insights_text += "</marketing_insights>\n"
# Create optimized document structure - truncate content to avoid token overload
content_preview = content[:500] + "..." if len(content) > 500 else content
formatted_docs.append(
f"""
<document id="{i}">
<title>{title}</title>
<date>{datetime}</date>
<contextual_summary>
{contextual_summary}
</contextual_summary>
{marketing_insights_text}
<content>
{content_preview}
</content>
</document>
"""
)
result = "\n".join(formatted_docs)
result = f"""
<search_results>
{result}
</search_results>
When using context from any document, reference the document title and date for attribution.
"""
return result
except Exception:
logger.opt(exception=True).debug("Error retrieving documents.")
return "Error retrieving documents."
@track(name="MongoDBRetrieverTool.parse_query")
def __parse_query(self, query: str) -> str:
try:
# Try to parse as JSON first
query_dict = json.loads(query)
return query_dict["query"]
except (json.JSONDecodeError, KeyError):
# If JSON parsing fails, return the query as-is
return query
|