mirror of
https://github.com/MODSetter/SurfSense.git
synced 2025-09-01 18:19:08 +00:00
243 lines
7.9 KiB
Python
243 lines
7.9 KiB
Python
from typing import Any, NamedTuple
|
|
|
|
from langchain_core.messages import BaseMessage
|
|
from litellm import get_model_info, token_counter
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class Section(BaseModel):
|
|
"""A section in the answer outline."""
|
|
|
|
section_id: int = Field(..., description="The zero-based index of the section")
|
|
section_title: str = Field(..., description="The title of the section")
|
|
questions: list[str] = Field(
|
|
..., description="Questions to research for this section"
|
|
)
|
|
|
|
|
|
class AnswerOutline(BaseModel):
|
|
"""The complete answer outline with all sections."""
|
|
|
|
answer_outline: list[Section] = Field(
|
|
..., description="List of sections in the answer outline"
|
|
)
|
|
|
|
|
|
class DocumentTokenInfo(NamedTuple):
|
|
"""Information about a document and its token cost."""
|
|
|
|
index: int
|
|
document: dict[str, Any]
|
|
formatted_content: str
|
|
token_count: int
|
|
|
|
|
|
def get_connector_emoji(connector_name: str) -> str:
|
|
"""Get an appropriate emoji for a connector type."""
|
|
connector_emojis = {
|
|
"YOUTUBE_VIDEO": "📹",
|
|
"EXTENSION": "🧩",
|
|
"CRAWLED_URL": "🌐",
|
|
"FILE": "📄",
|
|
"SLACK_CONNECTOR": "💬",
|
|
"NOTION_CONNECTOR": "📘",
|
|
"GITHUB_CONNECTOR": "🐙",
|
|
"LINEAR_CONNECTOR": "📊",
|
|
"JIRA_CONNECTOR": "🎫",
|
|
"DISCORD_CONNECTOR": "🗨️",
|
|
"TAVILY_API": "🔍",
|
|
"LINKUP_API": "🔗",
|
|
"GOOGLE_CALENDAR_CONNECTOR": "📅",
|
|
}
|
|
return connector_emojis.get(connector_name, "🔎")
|
|
|
|
|
|
def get_connector_friendly_name(connector_name: str) -> str:
|
|
"""Convert technical connector IDs to user-friendly names."""
|
|
connector_friendly_names = {
|
|
"YOUTUBE_VIDEO": "YouTube",
|
|
"EXTENSION": "Browser Extension",
|
|
"CRAWLED_URL": "Web Pages",
|
|
"FILE": "Files",
|
|
"SLACK_CONNECTOR": "Slack",
|
|
"NOTION_CONNECTOR": "Notion",
|
|
"GITHUB_CONNECTOR": "GitHub",
|
|
"LINEAR_CONNECTOR": "Linear",
|
|
"JIRA_CONNECTOR": "Jira",
|
|
"CONFLUENCE_CONNECTOR": "Confluence",
|
|
"GOOGLE_CALENDAR_CONNECTOR": "Google Calendar",
|
|
"DISCORD_CONNECTOR": "Discord",
|
|
"TAVILY_API": "Tavily Search",
|
|
"LINKUP_API": "Linkup Search",
|
|
}
|
|
return connector_friendly_names.get(connector_name, connector_name)
|
|
|
|
|
|
def convert_langchain_messages_to_dict(
|
|
messages: list[BaseMessage],
|
|
) -> list[dict[str, str]]:
|
|
"""Convert LangChain messages to format expected by token_counter."""
|
|
role_mapping = {"system": "system", "human": "user", "ai": "assistant"}
|
|
|
|
converted_messages = []
|
|
for msg in messages:
|
|
role = role_mapping.get(getattr(msg, "type", None), "user")
|
|
converted_messages.append({"role": role, "content": str(msg.content)})
|
|
|
|
return converted_messages
|
|
|
|
|
|
def format_document_for_citation(document: dict[str, Any]) -> str:
|
|
"""Format a single document for citation in the standard XML format."""
|
|
content = document.get("content", "")
|
|
doc_info = document.get("document", {})
|
|
document_id = document.get("chunk_id", "")
|
|
document_type = doc_info.get("document_type", "CRAWLED_URL")
|
|
|
|
return f"""<document>
|
|
<metadata>
|
|
<source_id>{document_id}</source_id>
|
|
<source_type>{document_type}</source_type>
|
|
</metadata>
|
|
<content>
|
|
{content}
|
|
</content>
|
|
</document>"""
|
|
|
|
|
|
def format_documents_section(
|
|
documents: list[dict[str, Any]], section_title: str = "Source material"
|
|
) -> str:
|
|
"""Format multiple documents into a complete documents section."""
|
|
if not documents:
|
|
return ""
|
|
|
|
formatted_docs = [format_document_for_citation(doc) for doc in documents]
|
|
|
|
return f"""{section_title}:
|
|
<documents>
|
|
{chr(10).join(formatted_docs)}
|
|
</documents>"""
|
|
|
|
|
|
def calculate_document_token_costs(
|
|
documents: list[dict[str, Any]], model: str
|
|
) -> list[DocumentTokenInfo]:
|
|
"""Pre-calculate token costs for each document."""
|
|
document_token_info = []
|
|
|
|
for i, doc in enumerate(documents):
|
|
formatted_doc = format_document_for_citation(doc)
|
|
|
|
# Calculate token count for this document
|
|
token_count = token_counter(
|
|
messages=[{"role": "user", "content": formatted_doc}], model=model
|
|
)
|
|
|
|
document_token_info.append(
|
|
DocumentTokenInfo(
|
|
index=i,
|
|
document=doc,
|
|
formatted_content=formatted_doc,
|
|
token_count=token_count,
|
|
)
|
|
)
|
|
|
|
return document_token_info
|
|
|
|
|
|
def find_optimal_documents_with_binary_search(
|
|
document_tokens: list[DocumentTokenInfo], available_tokens: int
|
|
) -> list[DocumentTokenInfo]:
|
|
"""Use binary search to find the maximum number of documents that fit within token limit."""
|
|
if not document_tokens or available_tokens <= 0:
|
|
return []
|
|
|
|
left, right = 0, len(document_tokens)
|
|
optimal_docs = []
|
|
|
|
while left <= right:
|
|
mid = (left + right) // 2
|
|
current_docs = document_tokens[:mid]
|
|
current_token_sum = sum(doc_info.token_count for doc_info in current_docs)
|
|
|
|
if current_token_sum <= available_tokens:
|
|
optimal_docs = current_docs
|
|
left = mid + 1
|
|
else:
|
|
right = mid - 1
|
|
|
|
return optimal_docs
|
|
|
|
|
|
def get_model_context_window(model_name: str) -> int:
|
|
"""Get the total context window size for a model (input + output tokens)."""
|
|
try:
|
|
model_info = get_model_info(model_name)
|
|
context_window = model_info.get("max_input_tokens", 4096) # Default fallback
|
|
return context_window
|
|
except Exception as e:
|
|
print(
|
|
f"Warning: Could not get model info for {model_name}, using default 4096 tokens. Error: {e}"
|
|
)
|
|
return 4096 # Conservative fallback
|
|
|
|
|
|
def optimize_documents_for_token_limit(
|
|
documents: list[dict[str, Any]], base_messages: list[BaseMessage], model_name: str
|
|
) -> tuple[list[dict[str, Any]], bool]:
|
|
"""
|
|
Optimize documents to fit within token limits using binary search.
|
|
|
|
Args:
|
|
documents: List of documents with content and metadata
|
|
base_messages: Base messages without documents (chat history + system + human message template)
|
|
model_name: Model name for token counting (required)
|
|
output_token_buffer: Number of tokens to reserve for model output
|
|
|
|
Returns:
|
|
Tuple of (optimized_documents, has_documents_remaining)
|
|
"""
|
|
if not documents:
|
|
return [], False
|
|
|
|
model = model_name
|
|
context_window = get_model_context_window(model)
|
|
|
|
# Calculate base token cost
|
|
base_messages_dict = convert_langchain_messages_to_dict(base_messages)
|
|
base_tokens = token_counter(messages=base_messages_dict, model=model)
|
|
available_tokens_for_docs = context_window - base_tokens
|
|
|
|
print(
|
|
f"Token optimization: Context window={context_window}, Base={base_tokens}, Available for docs={available_tokens_for_docs}"
|
|
)
|
|
|
|
if available_tokens_for_docs <= 0:
|
|
print("No tokens available for documents after base content and output buffer")
|
|
return [], False
|
|
|
|
# Calculate token costs for all documents
|
|
document_token_info = calculate_document_token_costs(documents, model)
|
|
|
|
# Find optimal number of documents using binary search
|
|
optimal_doc_info = find_optimal_documents_with_binary_search(
|
|
document_token_info, available_tokens_for_docs
|
|
)
|
|
|
|
# Extract the original document objects
|
|
optimized_documents = [doc_info.document for doc_info in optimal_doc_info]
|
|
has_documents_remaining = len(optimized_documents) > 0
|
|
|
|
print(
|
|
f"Token optimization result: Using {len(optimized_documents)}/{len(documents)} documents"
|
|
)
|
|
|
|
return optimized_documents, has_documents_remaining
|
|
|
|
|
|
def calculate_token_count(messages: list[BaseMessage], model_name: str) -> int:
|
|
"""Calculate token count for a list of LangChain messages."""
|
|
model = model_name
|
|
messages_dict = convert_langchain_messages_to_dict(messages)
|
|
return token_counter(messages=messages_dict, model=model)
|