import requests import json import openai from pydantic import BaseModel, Field from typing import Optional, Dict, Any, List, Union class Filter: """ 1. Analyzes the latest user messages with OpenAI to generate an optimized and context-based search query 2. Uses this optimized query for a search via Brave API (Pro AI subscription) 3. Adds the summary or search results to the LLM request """ class Valves(BaseModel): # OpenAI API settings OPENAI_API_KEY: str = Field(default="", description="Your OpenAI API key") OPENAI_MODEL: str = Field( default="gpt-3.5-turbo", description="OpenAI model to use" ) MAX_CHAT_HISTORY: int = Field( default=5, description="Number of past user messages to analyze" ) # Brave API settings BRAVE_API_KEY: str = Field(default="", description="Your Brave Search API key") INCLUDE_SOURCES: bool = Field( default=True, description="Include source URLs in the summary" ) NUM_FALLBACK_RESULTS: int = Field( default=3, description="Number of web results if no summary" ) MAX_SUMMARY_LENGTH: int = Field( default=2000, description="Maximum summary length" ) DEBUG_MODE: bool = Field(default=True, description="Enable debug output") # Format settings OUTPUT_MESSAGE: str = Field( default=" [Brave Summary + OpenAI Query]", description="Message to append to response", ) SUMMARY_PREFIX: str = Field( default="\n\nWeb search summary:\n", description="Prefix for search summary" ) RESULTS_PREFIX: str = Field( default="\n\nWeb search results:\n", description="Prefix for search results" ) def __init__(self): self.valves = self.Valves() self.debug_log = [] def _log_debug(self, message: str): if self.valves.DEBUG_MODE: self.debug_log.append(message) def _extract_recent_user_messages(self, messages: List[Dict]) -> List[Dict]: """Extract the most recent user messages from the chat history.""" user_messages = [] # Collect all user messages for message in messages: if message.get("role") == "user" and "content" in message: user_messages.append({"role": "user", "content": message["content"]}) # Return the most recent N messages (or all if fewer) max_messages = min(len(user_messages), self.valves.MAX_CHAT_HISTORY) return user_messages[-max_messages:] def _generate_optimized_query( self, user_messages: List[Dict], current_query: str ) -> str: """Use OpenAI to generate an optimized search query based on chat history.""" if not self.valves.OPENAI_API_KEY: self._log_debug("ERROR: No OpenAI API key provided, using original query") return current_query try: # Configure OpenAI client client = openai.OpenAI(api_key=self.valves.OPENAI_API_KEY) # Create prompt messages prompt_messages = [ { "role": "system", "content": """You are a search query optimizer. Your task is to analyze the conversation history and the current query, then generate an optimized search query that will provide the most relevant information. Focus on extracting key concepts and search terms. IMPORTANT: Return ONLY the optimized search query without any explanations or additional text.""", } ] # Add user message history prompt_messages.extend(user_messages[:-1]) # All but the current message # Add final instruction with current query prompt_messages.append( { "role": "user", "content": f"Based on our conversation history, create an optimized search query for: '{current_query}' ", } ) self._log_debug( f"Sending request to OpenAI with {len(prompt_messages)} messages" ) # Call the OpenAI API response = client.chat.completions.create( model=self.valves.OPENAI_MODEL, messages=prompt_messages, temperature=0.3, max_tokens=100, ) # Extract the optimized query optimized_query = response.choices[0].message.content.strip() self._log_debug(f"Original query: '{current_query}'") self._log_debug(f"Optimized query: '{optimized_query}'") return optimized_query except Exception as e: error_msg = f"Error generating optimized query: {str(e)}" self._log_debug(error_msg) return current_query # Fall back to the original query def _get_web_search_with_summary_key(self, query: str) -> Dict[str, Any]: """Perform a web search via Brave API and get a summary key.""" self._log_debug(f"Starting web search for query: '{query}'") if not self.valves.BRAVE_API_KEY: self._log_debug("ERROR: No Brave API key provided") return {"error": "No API key provided"} endpoint = "https://api.search.brave.com/res/v1/web/search" headers = { "Accept": "application/json", "Accept-Encoding": "gzip", "X-Subscription-Token": self.valves.BRAVE_API_KEY, "Api-Version": "2023-10-11", } params = { "q": query, "summary": 1, "count": 5, "country": "DE", "search_lang": "de", "safesearch": "off", } self._log_debug(f"Web search API request: {endpoint} with params {params}") try: response = requests.get(endpoint, headers=headers, params=params) status_code = response.status_code self._log_debug(f"Web search API response status: {status_code}") data = response.json() self._log_debug(f"Web search response keys: {list(data.keys())}") # Check if we have a summarizer key if "summarizer" in data and "key" in data["summarizer"]: self._log_debug(f"Found summary key: {data['summarizer']['key']}") else: self._log_debug("No summary key found in response") # Check if we have web results if "web" in data and "results" in data["web"]: result_count = len(data["web"]["results"]) self._log_debug(f"Found {result_count} web results") else: self._log_debug("No web results found in response") return data except Exception as e: error_msg = f"Error performing Brave web search: {str(e)}" self._log_debug(error_msg) return {"error": error_msg} def _get_summary(self, summary_key: str) -> Dict[str, Any]: """Get the summary using the key from the web search.""" self._log_debug(f"Starting summary request with key: {summary_key}") if not summary_key: return {"error": "No summary key provided"} endpoint = "https://api.search.brave.com/res/v1/summarizer/search" headers = { "Accept": "application/json", "Accept-Encoding": "gzip", "X-Subscription-Token": self.valves.BRAVE_API_KEY, "Api-Version": "2024-04-23", } params = {"key": summary_key, "entity_info": 1} self._log_debug(f"Summary API request: {endpoint} with key={summary_key}") try: response = requests.get(endpoint, headers=headers, params=params) status_code = response.status_code self._log_debug(f"Summary API response status: {status_code}") data = response.json() self._log_debug(f"Summary response keys: {list(data.keys())}") # Check summary format - could be array or dict if "summary" in data: summary_type = type(data["summary"]).__name__ self._log_debug(f"Summary object type: {summary_type}") # Check if summary is a list/array (token format) if isinstance(data["summary"], list) and len(data["summary"]) > 0: first_item = data["summary"][0] if isinstance(first_item, dict) and "data" in first_item: self._log_debug( f"Found summary in token format with data field" ) else: self._log_debug("No content field in summary object") return data except Exception as e: error_msg = f"Error getting Brave summary: {str(e)}" self._log_debug(error_msg) return {"error": error_msg} def _format_web_results( self, web_data: Dict[str, Any], num_results: int = 3 ) -> str: """Format regular web search results as fallback.""" if ( "web" not in web_data or "results" not in web_data["web"] or not web_data["web"]["results"] ): return f"{self.valves.RESULTS_PREFIX}No web results available." results = web_data["web"]["results"][:num_results] formatted_text = self.valves.RESULTS_PREFIX for i, result in enumerate(results): formatted_text += f"{i+1}. {result.get('title', 'No title')}\n" if "description" in result: formatted_text += f"{result['description']}\n" formatted_text += f"Source: {result.get('url', 'No URL')}\n\n" return formatted_text def _extract_summary_content(self, summary_obj: Union[List, Dict, Any]) -> str: """Extract summary content from various possible formats.""" # Handle list format with token type objects if isinstance(summary_obj, list) and len(summary_obj) > 0: # Try to get data from first token first_item = summary_obj[0] if isinstance(first_item, dict) and "data" in first_item: return first_item["data"] # If no data field, try to concatenate all items as strings try: return "".join(str(item) for item in summary_obj) except: pass # Handle dictionary format with content field elif isinstance(summary_obj, dict) and "content" in summary_obj: return summary_obj["content"] # Return empty string if no matching format return "" def _format_summary_content( self, summary_data: Dict[str, Any], web_data: Dict[str, Any] ) -> str: """Format summary data or fall back to web results if no summary available.""" self._log_debug("Formatting summary content") # Check for errors if "error" in summary_data: self._log_debug(f"Error found: {summary_data['error']}") return self._format_web_results(web_data, self.valves.NUM_FALLBACK_RESULTS) # Look for and extract summary content content = "" if "summary" in summary_data: # Try to extract content from the summary object content = self._extract_summary_content(summary_data["summary"]) # If we have content, format it if content: self._log_debug( f"Successfully extracted summary content ({len(content)} chars)" ) formatted_text = self.valves.SUMMARY_PREFIX # Limit summary length if needed if len(content) > self.valves.MAX_SUMMARY_LENGTH: content = content[: self.valves.MAX_SUMMARY_LENGTH] + "..." formatted_text += content # Add sources if available and requested if self.valves.INCLUDE_SOURCES: # Try to extract citations from different places citations = [] if "citations" in summary_data.get("summary", {}): citations = summary_data["summary"]["citations"] elif ( "enrichments" in summary_data and "citations" in summary_data["enrichments"] ): citations = summary_data["enrichments"]["citations"] if citations: formatted_text += "\n\nSources:\n" for i, citation in enumerate(citations): if "url" in citation: formatted_text += f"{i+1}. {citation.get('url')}\n" return formatted_text else: # Fall back to web results self._log_debug("No summary content available, falling back to web results") return self._format_web_results(web_data, self.valves.NUM_FALLBACK_RESULTS) def inlet(self, body: dict) -> dict: """Process user query with chat history analysis and Brave search.""" self.debug_log = [] # Reset debug log self._log_debug("=== Starting new enhanced search request ===") if "messages" in body and body["messages"]: # Check if latest message is from user latest_message = body["messages"][-1] if latest_message.get("role") == "user" and "content" in latest_message: original_query = latest_message["content"] if original_query.strip(): # Extract recent user messages for context user_messages = self._extract_recent_user_messages(body["messages"]) self._log_debug( f"Extracted {len(user_messages)} recent user messages for context" ) # Generate optimized query with OpenAI if we have an API key if self.valves.OPENAI_API_KEY: optimized_query = self._generate_optimized_query( user_messages, original_query ) search_query = optimized_query else: search_query = original_query self._log_debug( "No OpenAI API key provided, using original query" ) # Use the optimized query for Brave search web_search_data = self._get_web_search_with_summary_key( search_query ) if "error" in web_search_data: self._log_debug(f"Web search error: {web_search_data['error']}") formatted_results = f"{self.valves.RESULTS_PREFIX}Failed to perform search: {web_search_data['error']}" else: # Extract summary key summary_key = web_search_data.get("summarizer", {}).get("key") if summary_key: # Get and format summary summary_data = self._get_summary(summary_key) formatted_results = self._format_summary_content( summary_data, web_search_data ) else: # Fall back to regular web results self._log_debug( "No summary key available, falling back to web results" ) formatted_results = self._format_web_results( web_search_data, self.valves.NUM_FALLBACK_RESULTS ) # Add info about the optimized query if different from original if search_query != original_query: query_info = f'\n\nOptimized search query: "{search_query}"\n' formatted_results = query_info + formatted_results # Append results to the message self._log_debug("Appending search results to user message") latest_message["content"] += formatted_results return body def stream(self, event: dict) -> dict: """Pass through streaming events unchanged.""" return event def outlet(self, body: dict) -> dict: """Add debug info to assistant response.""" if "messages" in body and body["messages"]: # Find the most recent assistant message last_assistant_message = None for message in body["messages"]: if message.get("role") == "assistant" and "content" in message: last_assistant_message = message # Append info to the assistant's response if last_assistant_message: last_assistant_message["content"] += self.valves.OUTPUT_MESSAGE if self.valves.DEBUG_MODE and self.debug_log: debug_text = "\n\n=== DEBUG LOG ===\n" for i, log in enumerate(self.debug_log, 1): debug_text += f"{i}. {log}\n" last_assistant_message["content"] += debug_text return body