OverviewFilesCommitsIssuesDocsReleasesPeers

ollama_semantic_function_calling commits / 8d12855

getting opensearchng to work , also history / metrics / topics
8d12855d414f644fe5864ec9dff0f405c6527abf
by fleeky.ashatron <fleeky@7rnx.net> · 2025-04-20T01:11:51.000Z

Changed files (12) +2406 -41

ollamatools-new/matrix_cmd/matrix_client.py +152 -20
……@@ -28,9 +28,14 @@ logger = logging.getLogger(__name__)
2828
2929# Import from ollamatools Matrix integration
3030try:
31 from ollamatools.plugins.integrations.matrix.client import (
32 MatrixClient as OllamaToolsMatrixClient,
33 MatrixConfig as OllamaToolsMatrixConfig,
31 # Import from the new enhanced_matrix_core module
32 from ollamatools.tools.enhanced_matrix_core import (
33 EnhancedMatrixClient,
34 _connect_enhanced_matrix
35 )
36
37 # Import from nio for basic Matrix functionality
38 from nio import (
3439 AsyncClient,
3540 MatrixRoom,
3641 RoomMessageText,
……@@ -38,16 +43,106 @@ try:
3843 JoinError
3944 )
4045
41 # Re-export MatrixConfig
42 MatrixConfig = OllamaToolsMatrixConfig
43
44 from ollamatools.plugins.integrations.matrix.commands import (
45 connect_enhanced_matrix,
46 handle_history_command,
47 handle_topics_command,
48 handle_metrics_command
49 )
50
46 # Define MatrixConfig class
47 class MatrixConfig:
48 """
49 Configuration handler for Matrix integration.
50 """
51
52 @staticmethod
53 def load_config(config_path: str) -> Dict:
54 """
55 Load configuration from a YAML file.
56
57 Args:
58 config_path: Path to the configuration file
59
60 Returns:
61 Dict: The configuration as a dictionary
62 """
63 try:
64 with open(config_path, 'r') as f:
65 config = yaml.safe_load(f)
66
67 logger.info(f"Loaded configuration from {config_path}")
68 return config
69
70 except Exception as e:
71 logger.error(f"Error loading configuration from {config_path}: {e}")
72 raise
73
74 @staticmethod
75 def save_config(config: Dict, config_path: str) -> None:
76 """
77 Save configuration to a YAML file.
78
79 Args:
80 config: The configuration as a dictionary
81 config_path: Path to save the configuration to
82 """
83 try:
84 # Create directory if it doesn't exist
85 os.makedirs(os.path.dirname(config_path), exist_ok=True)
86
87 with open(config_path, 'w') as f:
88 yaml.dump(config, f, default_flow_style=False)
89
90 logger.info(f"Saved configuration to {config_path}")
91
92 except Exception as e:
93 logger.error(f"Error saving configuration to {config_path}: {e}")
94 raise
95
96 @staticmethod
97 def create_default_config() -> Dict:
98 """
99 Create a default configuration.
100
101 Returns:
102 Dict: The default configuration
103 """
104 return {
105 "MATRIX": {
106 "HOMESERVER": "https://matrix.example.org",
107 "USER": "@user:example.org",
108 "PASSWORD": "your-password",
109 "DEVICE_ID": "matrix-cmd-agent",
110 "ROOMS": [
111 "#room:example.org"
112 ],
113 "ADMIN_USERS": [
114 "@admin:example.org"
115 ]
116 },
117 "OLLAMA": {
118 "HOST": "localhost:11434",
119 "MODEL": "llama3:latest",
120 "PROMPT": "You are a helpful assistant.",
121 "AVAILABLE_MODELS": [
122 "llama3:latest",
123 "mistral:latest",
124 "gemma:latest"
125 ]
126 },
127 "ENHANCED": {
128 "ENABLE_SEARCH": False,
129 "ENABLE_CONVERSATION_TRACKING": False,
130 "HISTORY_FILE": "~/.config/matrix-cmd/history.json",
131 "MAX_CONVERSATION_HISTORY": 100,
132 "RESPOND_TO_ALL_MESSAGES": False,
133 "OPENSEARCH": {
134 "BASE_URL": "http://localhost:8081",
135 "ENGINES": [
136 "google",
137 "bing",
138 "duckduckgo"
139 ],
140 "NUM_RESULTS": 5
141 }
142 }
143 }
144
145 # Import functions from plugins
51146 from ollamatools.plugins.integrations.matrix.search import (
52147 handle_search_command,
53148 get_default_opensearch_config
……@@ -61,6 +156,24 @@ try:
61156 get_performance_metrics
62157 )
63158
159 # Define connect_enhanced_matrix function
160 connect_enhanced_matrix = _connect_enhanced_matrix
161
162 # Define handle_history_command function
163 async def handle_history_command(client, count, room_id):
164 from ollamatools.plugins.integrations.matrix.history import handle_history_command
165 await handle_history_command(client, count, room_id)
166
167 # Define handle_topics_command function
168 async def handle_topics_command(client, room_id):
169 from ollamatools.plugins.integrations.matrix.history import handle_topics_command
170 await handle_topics_command(client, room_id)
171
172 # Define handle_metrics_command function
173 async def handle_metrics_command(client, room_id):
174 from ollamatools.plugins.integrations.matrix.history import handle_metrics_command
175 await handle_metrics_command(client, room_id)
176
64177 OLLAMATOOLS_AVAILABLE = True
65178except ImportError:
66179 logger.error(
……@@ -1017,16 +1130,35 @@ For example: "Hey {username}, what's the weather like today?"
10171130 await self.send_message(room_id, f"Searching for: {query}...")
10181131
10191132 try:
1020 # Simple implementation that uses Ollama to generate a response
1021 # In a real implementation, this would use OpenSearchNG
1022 search_prompt = f"Please provide a brief search result for the query: {query}"
1023 response = await self.process_with_ollama(search_prompt, None, room_id)
1133 # Use OpenSearchNG to get real search results
1134 from ollamatools.plugins.integrations.opensearch.tools import opensearch_query
1135
1136 # Execute the search query
1137 results = opensearch_query(query, self.opensearch_config)
1138
1139 if not results or not results.get("results"):
1140 await self.send_message(room_id, f"No results found for: {query}")
1141 return
10241142
1025 # Format the response
1026 formatted_response = f"Search Results for '{query}':\n\n{response}"
1143 # Format the results
1144 formatted_results = ["Search Results:"]
1145
1146 for i, result in enumerate(results["results"][:5], 1):
1147 title = result.get("title", "No title")
1148 url = result.get("url", "No URL")
1149 content = result.get("content", "No content")
1150
1151 # Truncate content if it's too long
1152 if len(content) > 200:
1153 content = content[:200] + "..."
1154
1155 formatted_results.append(f"{i}. {title}")
1156 formatted_results.append(f" URL: {url}")
1157 formatted_results.append(f" {content}")
1158 formatted_results.append("")
10271159
10281160 # Send the formatted response
1029 await self.send_message(room_id, formatted_response)
1161 await self.send_message(room_id, "\n".join(formatted_results))
10301162
10311163 except Exception as e:
10321164 logger.error(f"Error executing search: {e}")
ollamatools-new/ollamatools/__init__.py +2 -19
……@@ -2,26 +2,9 @@
22
33from ollamatools.version import __version__
44
5# Expose core components
6from ollamatools.core.agent import Agent
7from ollamatools.core.config import Config
8
9# Expose agent implementations
10from ollamatools.agents.basic import BasicAgent
11from ollamatools.agents.modular import ModularAgent
12from ollamatools.agents.extended import ExtendedAgent
13
14# Expose plugin system
15from ollamatools.plugins.base import Plugin
16from ollamatools.plugins.registry import PluginRegistry
5# Only expose the version in __init__.py to avoid circular imports
6# Other components should be imported directly from their modules
177
188__all__ = [
199 "__version__",
20 "Agent",
21 "Config",
22 "BasicAgent",
23 "ModularAgent",
24 "ExtendedAgent",
25 "Plugin",
26 "PluginRegistry",
2710]
ollamatools-new/ollamatools/agents/basic.py +2 -2
……@@ -11,10 +11,10 @@ from typing import Dict, Any, List, Optional
1111
1212from ollamatools.core.agent import OllamaAgent
1313from ollamatools.core.config import Config
14from ollamatools.core.utils.logger import setup_logger
14from ollamatools.core.utils.logger import get_ollamatools_logger
1515
1616# Set up logging
17logger = setup_logger("basic_ollama_agent", level="INFO")
17logger = get_ollamatools_logger("basic_ollama_agent", level="INFO")
1818
1919class BasicOllamaAgent:
2020 """
ollamatools-new/ollamatools/tools/__init__.py +15 -0
……@@ -0,0 +1,15 @@
1"""
2Tools module for OllamaTools.
3
4This module provides various utility tools for OllamaTools.
5"""
6
7from ollamatools.tools.enhanced_matrix_core import (
8 EnhancedMatrixClient,
9 _connect_enhanced_matrix
10)
11
12__all__ = [
13 'EnhancedMatrixClient',
14 '_connect_enhanced_matrix'
15]
ollamatools-new/ollamatools/tools/enhanced_matrix_core.py +736 -0
……@@ -0,0 +1,736 @@
1"""
2Enhanced Matrix core functionality for OllamaTools.
3
4This module provides the core functionality for the Enhanced Matrix integration,
5including the EnhancedMatrixClient class and connection functions.
6
7Dependencies:
8- matrix-nio: For Matrix protocol support
9- matrix-nio[e2e]: For end-to-end encryption support (optional)
10- pyyaml: For configuration handling
11- ollama: For Ollama API access
12"""
13
14import os
15import sys
16import yaml
17import logging
18import asyncio
19import re
20import time
21from typing import Dict, List, Optional, Union, Callable, Any, Set
22from collections import defaultdict
23
24# Configure logging
25logger = logging.getLogger(__name__)
26
27# Import from matrix-nio
28try:
29 from nio import (
30 AsyncClient,
31 MatrixRoom,
32 RoomMessageText,
33 LoginResponse,
34 JoinError
35 )
36except ImportError:
37 raise ImportError(
38 "Matrix support requires the 'matrix-nio' package. "
39 "Please install it with 'pip install matrix-nio' or "
40 "'pip install matrix-nio[e2e]' for end-to-end encryption support."
41 )
42
43# Make ollama import optional
44try:
45 from ollama import AsyncClient as OllamaAsync
46 OLLAMA_AVAILABLE = True
47except ImportError:
48 OLLAMA_AVAILABLE = False
49 logger.warning(
50 "The 'ollama' package is not installed. "
51 "Message processing with Ollama will not be available. "
52 "Install it with 'pip install ollama'."
53 )
54
55class EnhancedMatrixClient:
56 """
57 Enhanced client for connecting to Matrix homeservers and interacting with rooms.
58
59 This client extends the basic Matrix client with advanced features like:
60 - Search integration
61 - Conversation history tracking
62 - Topic analysis
63 - Performance metrics
64 """
65
66 # Command patterns
67 MODEL_COMMAND_PATTERN = re.compile(r'^/model\s+([a-zA-Z0-9:_\-\.]+)$')
68 HELP_COMMAND_PATTERN = re.compile(r'^/help$')
69 LIST_MODELS_COMMAND_PATTERN = re.compile(r'^/models$')
70 PING_COMMAND_PATTERN = re.compile(r'^/ping$')
71 SEARCH_COMMAND_PATTERN = re.compile(r'^/search\s+(.+)$')
72 HISTORY_COMMAND_PATTERN = re.compile(r'^/history(?:\s+(\d+))?$')
73 TOPICS_COMMAND_PATTERN = re.compile(r'^/topics$')
74 METRICS_COMMAND_PATTERN = re.compile(r'^/metrics$')
75
76 def __init__(
77 self,
78 homeserver: str,
79 user_id: str,
80 password: str,
81 device_id: str = None,
82 ollama_host: str = "localhost:11434",
83 ollama_model: str = "llama3:latest",
84 ollama_prompt: str = None,
85 available_models: List[str] = None,
86 admin_users: List[str] = None,
87 opensearch_config: Dict[str, Any] = None,
88 ssl: bool = True,
89 enable_search: bool = False,
90 enable_conversation_tracking: bool = False,
91 history_file: str = None,
92 max_context_length: int = 10,
93 max_conversation_history: int = 100,
94 respond_to_all_messages: bool = False
95 ):
96 """
97 Initialize a new Enhanced Matrix client.
98
99 Args:
100 homeserver: The URL of the Matrix homeserver
101 user_id: The Matrix user ID (e.g., @user:example.com)
102 password: The password for the Matrix account
103 device_id: A unique identifier for this device (for E2E encryption)
104 ollama_host: The host address for the Ollama API
105 ollama_model: The Ollama model to use for generating responses
106 ollama_prompt: An optional system prompt for the Ollama model
107 available_models: List of available Ollama models
108 admin_users: List of Matrix user IDs that can change models for all users
109 opensearch_config: Configuration for OpenSearchNG
110 ssl: Whether to use SSL for the Matrix connection
111 enable_search: Whether to enable search functionality
112 enable_conversation_tracking: Whether to enable conversation history tracking
113 history_file: Path to the conversation history file
114 max_context_length: Maximum number of messages to keep in context
115 max_conversation_history: Maximum number of conversations to keep in history
116 respond_to_all_messages: Whether to respond to all messages (not just mentions)
117 """
118 self.homeserver = homeserver
119 self.user_id = user_id
120 self.password = password
121 self.device_id = device_id
122 self.ollama_host = ollama_host
123 self.ollama_model = ollama_model
124 self.ollama_prompt = ollama_prompt
125 self.ssl = ssl
126
127 # Enhanced features
128 self.enable_search = enable_search
129 self.enable_conversation_tracking = enable_conversation_tracking
130 self.history_file = history_file or os.path.expanduser("~/.config/matrix-cmd/history.json")
131 self.max_context_length = max_context_length
132 self.max_conversation_history = max_conversation_history
133 self.respond_to_all_messages = respond_to_all_messages
134
135 # Set OpenSearch config
136 self.opensearch_config = opensearch_config or {}
137
138 # Conversation history
139 self.conversation_history = []
140
141 # Initialize the Matrix client
142 self.client = AsyncClient(
143 homeserver=self.homeserver,
144 device_id=self.device_id,
145 ssl=self.ssl
146 )
147
148 # Track when the client connected to avoid processing old messages
149 self.connection_time = None
150
151 # Track joined rooms
152 self.rooms = []
153
154 # Message callback
155 self.message_callback = None
156
157 # User context tracking
158 self.user_contexts = defaultdict(list)
159
160 # User model preferences
161 self.user_models = {}
162
163 # Available models
164 self.available_models = available_models or [ollama_model]
165
166 # Admin users who can change models
167 self.admin_users = set(admin_users or [])
168
169 async def connect(self, max_retries: int = 3, initial_retry_delay: float = 1.0) -> bool:
170 """
171 Connect to the Matrix homeserver and log in.
172
173 Args:
174 max_retries: Maximum number of retry attempts (default: 3)
175 initial_retry_delay: Initial delay between retries in seconds (default: 1.0)
176
177 Returns:
178 bool: True if the connection was successful, False otherwise
179 """
180 retry_count = 0
181 current_retry_delay = initial_retry_delay
182
183 while retry_count <= max_retries:
184 try:
185 # Extract username from user_id (e.g., "@user:example.com" -> "user")
186 username = self.user_id.split(':')[0].lstrip('@')
187
188 # Set the user ID for the client
189 self.client.user = username
190
191 logger.info(f"Attempting to connect to {self.homeserver} as {self.user_id}...")
192
193 response = await self.client.login(
194 password=self.password
195 )
196
197 if isinstance(response, LoginResponse):
198 logger.info(f"Successfully connected to {self.homeserver} as {self.user_id}")
199 self.connection_time = asyncio.get_event_loop().time()
200
201 # Load conversation history if tracking is enabled
202 if self.enable_conversation_tracking:
203 from ollamatools.plugins.integrations.matrix.history import load_conversation_history
204 self.conversation_history = load_conversation_history(self.history_file)
205
206 return True
207 else:
208 logger.error(f"Failed to connect to {self.homeserver}: {response}")
209
210 # Check if we should retry
211 if retry_count < max_retries:
212 retry_count += 1
213 logger.info(f"Retrying connection (attempt {retry_count}/{max_retries}) in {current_retry_delay:.1f} seconds...")
214 await asyncio.sleep(current_retry_delay)
215 # Exponential backoff
216 current_retry_delay *= 2
217 else:
218 logger.error(f"Maximum retry attempts ({max_retries}) reached. Giving up.")
219 return False
220
221 except asyncio.CancelledError:
222 logger.info("Connection attempt cancelled by user.")
223 raise
224
225 except Exception as e:
226 logger.error(f"Error connecting to Matrix server: {e}")
227
228 # Check if this is a rate limiting error
229 if "429" in str(e) and "ratelimited" in str(e):
230 # Try to extract the retry time from the error message
231 import re
232 retry_match = re.search(r'sleeping for (\d+)ms', str(e))
233
234 if retry_match:
235 retry_ms = int(retry_match.group(1))
236 retry_sec = retry_ms / 1000
237
238 if retry_sec > 300: # If wait time is more than 5 minutes
239 logger.error(f"Rate limit wait time is too long ({retry_sec:.1f} seconds). Consider trying again later.")
240 return False
241
242 logger.info(f"Server is rate limiting requests. Waiting for {retry_sec:.1f} seconds before retrying...")
243
244 try:
245 # Wait for the specified time
246 await asyncio.sleep(retry_sec)
247 # Don't increment retry count for rate limiting
248 continue
249 except asyncio.CancelledError:
250 logger.info("Rate limit wait cancelled by user.")
251 raise
252 else:
253 # If we can't extract the retry time, use exponential backoff
254 if retry_count < max_retries:
255 retry_count += 1
256 logger.info(f"Rate limited but couldn't determine wait time. Retrying (attempt {retry_count}/{max_retries}) in {current_retry_delay:.1f} seconds...")
257 await asyncio.sleep(current_retry_delay)
258 # Exponential backoff
259 current_retry_delay *= 2
260 else:
261 logger.error(f"Maximum retry attempts ({max_retries}) reached. Giving up.")
262 return False
263 else:
264 # For non-rate-limiting errors
265 if retry_count < max_retries:
266 retry_count += 1
267 logger.info(f"Retrying connection (attempt {retry_count}/{max_retries}) in {current_retry_delay:.1f} seconds...")
268 await asyncio.sleep(current_retry_delay)
269 # Exponential backoff
270 current_retry_delay *= 2
271 else:
272 logger.error(f"Maximum retry attempts ({max_retries}) reached. Giving up.")
273 return False
274
275 return False
276
277 async def join_room(self, room_id: str) -> bool:
278 """
279 Join a Matrix room.
280
281 Args:
282 room_id: The ID of the room to join (e.g., #room:example.com)
283
284 Returns:
285 bool: True if joining was successful, False otherwise
286 """
287 try:
288 response = await self.client.join(room_id)
289
290 if isinstance(response, JoinError):
291 logger.error(f"Failed to join room {room_id}: {response}")
292 return False
293
294 logger.info(f"Successfully joined room {room_id}")
295 self.rooms.append(room_id)
296 return True
297
298 except Exception as e:
299 logger.error(f"Error joining room {room_id}: {e}")
300 return False
301
302 async def send_message(self, room_id: str, message: str) -> bool:
303 """
304 Send a message to a Matrix room.
305
306 Args:
307 room_id: The ID of the room to send the message to
308 message: The message to send
309
310 Returns:
311 bool: True if sending was successful, False otherwise
312 """
313 try:
314 response = await self.client.room_send(
315 room_id=room_id,
316 message_type="m.room.message",
317 content={
318 "msgtype": "m.text",
319 "body": message
320 }
321 )
322
323 logger.info(f"Message sent successfully to room {room_id}")
324 return True
325
326 except Exception as e:
327 logger.error(f"Error sending message to room {room_id}: {e}")
328 return False
329
330 async def process_with_ollama(self, message: str, sender: str = None, room_id: str = None) -> str:
331 """
332 Process a message with Ollama and return the response.
333
334 Args:
335 message: The message to process
336 sender: The Matrix user ID of the sender
337 room_id: The ID of the room where the message was sent
338
339 Returns:
340 str: The response from Ollama
341 """
342 if not OLLAMA_AVAILABLE:
343 return "Ollama is not available. Please install the 'ollama' package with 'pip install ollama'."
344
345 try:
346 ollama_client = OllamaAsync(host=self.ollama_host)
347
348 # Determine which model to use
349 model = self.ollama_model
350 if sender and sender in self.user_models:
351 model = self.user_models[sender]
352
353 messages = []
354
355 # Add system prompt if provided
356 if self.ollama_prompt:
357 messages.append({
358 "role": "system",
359 "content": self.ollama_prompt
360 })
361
362 # Add conversation context if available
363 if sender and room_id:
364 context_key = f"{sender}:{room_id}"
365 for ctx_msg in self.user_contexts[context_key]:
366 messages.append(ctx_msg)
367
368 # Add user message
369 user_message = {
370 "role": "user",
371 "content": message
372 }
373 messages.append(user_message)
374
375 # Store user message in context
376 if sender and room_id:
377 context_key = f"{sender}:{room_id}"
378 self.user_contexts[context_key].append(user_message)
379 # Trim context if it gets too long
380 if len(self.user_contexts[context_key]) > self.max_context_length:
381 self.user_contexts[context_key] = self.user_contexts[context_key][-self.max_context_length:]
382
383 logger.info(f"Sending request to Ollama with model: {model}")
384 start_time = time.time()
385 response = await ollama_client.chat(
386 model=model,
387 messages=messages
388 )
389 response_time = time.time() - start_time
390
391 logger.info(f"Received response from Ollama: {response}")
392 response_message = response["message"]
393 message_content = response_message["content"]
394 logger.info(f"Extracted message content: {message_content}")
395
396 # Store assistant response in context
397 if sender and room_id:
398 context_key = f"{sender}:{room_id}"
399 self.user_contexts[context_key].append({
400 "role": "assistant",
401 "content": message_content
402 })
403 # Trim context if it gets too long
404 if len(self.user_contexts[context_key]) > self.max_context_length:
405 self.user_contexts[context_key] = self.user_contexts[context_key][-self.max_context_length:]
406
407 # Record conversation if tracking is enabled
408 if self.enable_conversation_tracking and sender and room_id:
409 from ollamatools.plugins.integrations.matrix.history import record_conversation
410 record_conversation(
411 self.conversation_history,
412 message,
413 message_content,
414 sender,
415 room_id,
416 response_time,
417 self.history_file,
418 self.max_conversation_history
419 )
420
421 return message_content
422
423 except Exception as e:
424 logger.error(f"Error processing message with Ollama: {e}")
425 return f"Error processing message: {e}"
426
427 async def _message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None:
428 """
429 Internal callback for handling incoming messages.
430
431 Args:
432 room: The room the message was sent in
433 event: The message event
434 """
435 if not isinstance(event, RoomMessageText):
436 return
437
438 # Get message timestamp
439 message_timestamp = event.server_timestamp / 1000
440
441 # Skip messages from before the client connected
442 if self.connection_time and message_timestamp < self.connection_time:
443 logger.debug("Skipping old message")
444 return
445
446 # Skip messages from the bot itself
447 if event.sender == self.user_id:
448 logger.debug("Skipping message from self")
449 return
450
451 message = event.body
452 room_id = room.room_id
453 sender = event.sender
454
455 logger.info(f"Received message in room {room_id} from {sender}: {message}")
456
457 # If a custom callback is set, use it
458 if self.message_callback:
459 await self.message_callback(room, event, self)
460 return
461
462 # Check if this is a command
463 if message.startswith('/'):
464 await self._handle_command(message, room_id, sender)
465 return
466
467 # Extract username from user_id (e.g., "@user:example.com" -> "user")
468 username = self.user_id.split(':')[0].lstrip('@')
469
470 # Check if the message contains the bot's username (case insensitive)
471 if username.lower() in message.lower() or self.respond_to_all_messages:
472 logger.info(f"Username mention detected in message: {message}")
473
474 # Check if this is a factual query and search is enabled
475 search_context = ""
476 if self.enable_search:
477 try:
478 from ollamatools.plugins.integrations.matrix.search import execute_search_for_message
479 search_context = await execute_search_for_message(message, self.opensearch_config)
480 except ImportError:
481 logger.warning("Search module not available")
482
483 # Process with Ollama and send response
484 if search_context:
485 # Add search context to the message
486 enhanced_message = f"{message}\n\nSearch context: {search_context}"
487 response = await self.process_with_ollama(enhanced_message, sender, room_id)
488 else:
489 response = await self.process_with_ollama(message, sender, room_id)
490
491 await self.send_message(room_id, response)
492 return
493
494 # For other messages without username mention, don't automatically respond
495
496 async def _handle_command(self, message: str, room_id: str, sender: str) -> None:
497 """
498 Handle commands sent to the bot.
499
500 Args:
501 message: The command message
502 room_id: The ID of the room where the command was sent
503 sender: The Matrix user ID of the sender
504 """
505 # Check if this is a model change command
506 model_match = self.MODEL_COMMAND_PATTERN.match(message)
507 if model_match:
508 model_name = model_match.group(1)
509 await self._change_model(model_name, room_id, sender)
510 return
511
512 # Check if this is a help command
513 if self.HELP_COMMAND_PATTERN.match(message):
514 await self._send_help(room_id)
515 return
516
517 # Check if this is a list models command
518 if self.LIST_MODELS_COMMAND_PATTERN.match(message):
519 await self._list_models(room_id)
520 return
521
522 # Check if this is a ping command
523 if self.PING_COMMAND_PATTERN.match(message):
524 await self.send_message(room_id, "Pong!")
525 return
526
527 # Check if this is a search command
528 search_match = self.SEARCH_COMMAND_PATTERN.match(message)
529 if search_match:
530 if not self.enable_search:
531 await self.send_message(room_id, "Search functionality is not enabled.")
532 return
533 query = search_match.group(1)
534
535 from ollamatools.plugins.integrations.matrix.search import handle_search_command
536 await handle_search_command(self, query, room_id)
537 return
538
539 # Check if this is a history command
540 history_match = self.HISTORY_COMMAND_PATTERN.match(message)
541 if history_match:
542 if not self.enable_conversation_tracking:
543 await self.send_message(room_id, "Conversation tracking is not enabled.")
544 return
545 count = int(history_match.group(1)) if history_match.group(1) else 10
546
547 from ollamatools.plugins.integrations.matrix.history import handle_history_command
548 await handle_history_command(self, count, room_id)
549 return
550
551 # Check if this is a topics command
552 if self.TOPICS_COMMAND_PATTERN.match(message):
553 if not self.enable_conversation_tracking:
554 await self.send_message(room_id, "Conversation tracking is not enabled.")
555 return
556
557 from ollamatools.plugins.integrations.matrix.history import handle_topics_command
558 await handle_topics_command(self, room_id)
559 return
560
561 # Check if this is a metrics command
562 if self.METRICS_COMMAND_PATTERN.match(message):
563 if not self.enable_conversation_tracking:
564 await self.send_message(room_id, "Conversation tracking is not enabled.")
565 return
566
567 from ollamatools.plugins.integrations.matrix.history import handle_metrics_command
568 await handle_metrics_command(self, room_id)
569 return
570
571 # Unknown command
572 await self.send_message(room_id, f"Unknown command: {message}\nType /help for available commands.")
573
574 async def _change_model(self, model_name: str, room_id: str, sender: str) -> None:
575 """
576 Change the Ollama model for a user.
577
578 Args:
579 model_name: The name of the model to change to
580 room_id: The ID of the room where the command was sent
581 sender: The Matrix user ID of the sender
582 """
583 # Check if the model is available
584 if model_name not in self.available_models:
585 await self.send_message(
586 room_id,
587 f"Model '{model_name}' is not available. Available models: {', '.join(self.available_models)}"
588 )
589 return
590
591 # Change the model for the user
592 self.user_models[sender] = model_name
593
594 # Clear the user's context
595 for key in list(self.user_contexts.keys()):
596 if key.startswith(f"{sender}:"):
597 self.user_contexts[key] = []
598
599 await self.send_message(room_id, f"Changed model to '{model_name}' for {sender}")
600 logger.info(f"Changed model to '{model_name}' for user {sender}")
601
602 async def _send_help(self, room_id: str) -> None:
603 """
604 Send help information to a room.
605
606 Args:
607 room_id: The ID of the room to send the help to
608 """
609 # Build help text based on enabled features
610 commands = [
611 "- /help - Show this help message",
612 "- /models - List available models",
613 "- /model <model_name> - Change to a different model",
614 "- /ping - Test if the bot is responsive"
615 ]
616
617 # Add search command if enabled
618 if self.enable_search:
619 commands.append("- /search <query> - Search the web for information")
620
621 # Add conversation tracking commands if enabled
622 if self.enable_conversation_tracking:
623 commands.append("- /history [count] - Show recent conversation history")
624 commands.append("- /topics - Show conversation topic analysis")
625 commands.append("- /metrics - Show performance metrics")
626
627 help_text = """
628Available commands:
629{commands}
630
631You can also mention me by my username in your message and I'll respond.
632For example: "Hey {username}, what's the weather like today?"
633 """.format(
634 commands="\n".join(commands),
635 username=self.user_id.split(':')[0].lstrip('@')
636 )
637 await self.send_message(room_id, help_text.strip())
638
639 async def _list_models(self, room_id: str) -> None:
640 """
641 List available models to a room.
642
643 Args:
644 room_id: The ID of the room to send the list to
645 """
646 models_text = f"Available models: {', '.join(self.available_models)}"
647 await self.send_message(room_id, models_text)
648
649 async def listen(self) -> None:
650 """
651 Start listening for messages in joined rooms.
652
653 This method will block until the client is stopped.
654 """
655 # Set up the message callback
656 self.client.add_event_callback(self._message_callback, RoomMessageText)
657
658 # Start syncing with the server
659 await self.client.sync_forever(timeout=30000, full_state=True)
660
661 async def disconnect(self) -> None:
662 """
663 Disconnect from the Matrix homeserver.
664 """
665 await self.client.logout()
666 await self.client.close()
667 logger.info(f"Disconnected from {self.homeserver}")
668
669
670async def _connect_enhanced_matrix(
671 homeserver: str,
672 user_id: str,
673 password: str,
674 device_id: str = None,
675 ollama_host: str = "localhost:11434",
676 ollama_model: str = "llama3:latest",
677 ollama_prompt: str = None,
678 available_models: List[str] = None,
679 admin_users: List[str] = None,
680 opensearch_config: Dict[str, Any] = None,
681 ssl: bool = True,
682 enable_search: bool = True,
683 enable_conversation_tracking: bool = True,
684 max_context_length: int = 10,
685 max_conversation_history: int = 100,
686 respond_to_all_messages: bool = False
687) -> Optional[EnhancedMatrixClient]:
688 """
689 Connect to a Matrix homeserver with enhanced features.
690
691 Args:
692 homeserver: The URL of the Matrix homeserver
693 user_id: The Matrix user ID (e.g., @user:example.com)
694 password: The password for the Matrix account
695 device_id: A unique identifier for this device (for E2E encryption)
696 ollama_host: The host address for the Ollama API
697 ollama_model: The Ollama model to use for generating responses
698 ollama_prompt: An optional system prompt for the Ollama model
699 available_models: List of available Ollama models
700 admin_users: List of Matrix user IDs that can change models for all users
701 opensearch_config: Configuration for OpenSearchNG
702 ssl: Whether to use SSL for the Matrix connection
703 enable_search: Whether to enable search functionality
704 enable_conversation_tracking: Whether to enable conversation tracking
705 max_context_length: Maximum number of messages to keep in context
706 max_conversation_history: Maximum number of conversations to keep in history
707 respond_to_all_messages: Whether to respond to all messages (not just mentions)
708
709 Returns:
710 Optional[EnhancedMatrixClient]: A connected Enhanced Matrix client, or None if connection failed
711 """
712 client = EnhancedMatrixClient(
713 homeserver=homeserver,
714 user_id=user_id,
715 password=password,
716 device_id=device_id,
717 ollama_host=ollama_host,
718 ollama_model=ollama_model,
719 ollama_prompt=ollama_prompt,
720 available_models=available_models,
721 admin_users=admin_users,
722 opensearch_config=opensearch_config,
723 ssl=ssl,
724 enable_search=enable_search,
725 enable_conversation_tracking=enable_conversation_tracking,
726 max_context_length=max_context_length,
727 max_conversation_history=max_conversation_history,
728 respond_to_all_messages=respond_to_all_messages
729 )
730
731 success = await client.connect()
732
733 if success:
734 return client
735 else:
736 return None
ollamatools-new/test_direct_opensearch.py +327 -0
……@@ -0,0 +1,327 @@
1#!/usr/bin/env python3
2"""
3Test script to directly verify OpenSearchNG functionality.
4"""
5
6import requests
7import json
8import logging
9import sys
10from typing import Dict, Any
11
12# Set up logging
13logging.basicConfig(level=logging.INFO)
14logger = logging.getLogger(__name__)
15
16# OpenSearchNG configuration
17OPENSEARCH_CONFIG = {
18 "base_url": "http://localhost:8081", # Default local OpenSearchNG URL
19 "search_path": "/search", # Path for search endpoint
20 "preferences_path": "/preferences", # Path for preferences endpoint
21 "categories": ["general"], # Default search categories
22 "engines": ["google", "bing", "duckduckgo"], # Default search engines
23 "language": "en-US", # Default language
24 "format": "json", # Response format (html, json, csv, rss)
25 "timeout": 10, # Default timeout in seconds
26 "num_results": 10, # Default number of results
27 "safe_search": 0 # 0=off, 1=moderate, 2=strict
28}
29
30def opensearch_query(query: str, config: Dict[str, Any] = None) -> Dict[str, Any]:
31 """
32 Perform a search query using a local OpenSearchNG instance.
33
34 Args:
35 query: Search query string
36 config: OpenSearchNG configuration (defaults used if None)
37
38 Returns:
39 dict: Search results with metadata
40 """
41 # Validate input
42 if not query or not isinstance(query, str):
43 raise ValueError("Query must be a non-empty string")
44
45 # Use default config if none provided
46 search_config = OPENSEARCH_CONFIG.copy()
47 if config:
48 search_config.update(config)
49
50 # Construct the search URL
51 base_url = search_config["base_url"]
52 search_path = search_config["search_path"]
53 search_url = f"{base_url}{search_path}"
54
55 # Prepare query parameters
56 params = {
57 "q": query,
58 "format": search_config["format"],
59 "categories": ",".join(search_config["categories"]),
60 "engines": ",".join(search_config["engines"]),
61 "language": search_config["language"],
62 "safesearch": search_config["safe_search"],
63 "time_range": "", # Optional time range
64 "pageno": 1, # First page of results
65 }
66
67 # Create session with appropriate headers
68 with requests.Session() as session:
69 headers = {
70 "User-Agent": "OllamaTools/0.2.0",
71 "Accept": "application/json"
72 }
73 session.headers.update(headers)
74
75 try:
76 logger.info(f"Making OpenSearchNG query: {query}")
77 response = session.get(
78 search_url,
79 params=params,
80 timeout=search_config["timeout"]
81 )
82
83 # Check for successful response
84 if response.status_code != 200:
85 logger.error(f"OpenSearchNG query failed with status code: {response.status_code}")
86 return {"error": f"Status code: {response.status_code}", "results": []}
87
88 # Parse the response based on format
89 if search_config["format"] == "json":
90 try:
91 result_data = response.json()
92 except json.JSONDecodeError:
93 logger.error("Failed to parse JSON response from OpenSearchNG")
94 return {"error": "JSON decode error", "results": []}
95 else:
96 # For other formats, return the raw content
97 result_data = {"raw_content": response.text}
98
99 # Process the results
100 processed_results = process_opensearch_results(result_data, query)
101
102 return processed_results
103
104 except requests.RequestException as e:
105 logger.error(f"OpenSearchNG request failed: {str(e)}")
106 return {"error": str(e), "results": []}
107
108def process_opensearch_results(result_data: Dict[str, Any], query: str) -> Dict[str, Any]:
109 """
110 Process and format the OpenSearchNG results.
111
112 Args:
113 result_data: Raw result data from OpenSearchNG
114 query: Original search query
115
116 Returns:
117 dict: Processed search results with metadata
118 """
119 processed_results = {
120 "query": query,
121 "results": [],
122 "suggestion": None,
123 "answer": None,
124 "infobox": None,
125 "engines": [],
126 "timing": {}
127 }
128
129 # Extract search results
130 if "results" in result_data:
131 for result in result_data["results"]:
132 processed_result = {
133 "title": result.get("title", ""),
134 "url": result.get("url", ""),
135 "content": result.get("content", ""),
136 "engine": result.get("engine", ""),
137 "score": result.get("score", 0),
138 "category": result.get("category", ""),
139 "pretty_url": result.get("pretty_url", "")
140 }
141 processed_results["results"].append(processed_result)
142
143 # Extract suggestions
144 if "suggestions" in result_data and result_data["suggestions"]:
145 try:
146 suggestion = result_data["suggestions"][0]
147 # Handle case where suggestion is a dict instead of a string
148 if isinstance(suggestion, dict):
149 processed_results["suggestion"] = suggestion.get("value", str(suggestion))
150 else:
151 processed_results["suggestion"] = suggestion
152 except (IndexError, TypeError) as e:
153 logger.warning(f"Error processing suggestions: {str(e)}")
154 processed_results["suggestion"] = None
155
156 # Extract answer if available
157 if "answers" in result_data and result_data["answers"]:
158 try:
159 answer = result_data["answers"][0]
160 # Handle case where answer is a dict instead of a string
161 if isinstance(answer, dict):
162 processed_results["answer"] = answer.get("content", str(answer))
163 else:
164 processed_results["answer"] = answer
165 except (IndexError, TypeError) as e:
166 logger.warning(f"Error processing answers: {str(e)}")
167 processed_results["answer"] = None
168
169 # Extract infobox if available
170 if "infoboxes" in result_data and result_data["infoboxes"]:
171 try:
172 infobox = result_data["infoboxes"][0]
173 # Handle case where infobox is a dict (which is expected)
174 processed_results["infobox"] = infobox
175 except (IndexError, TypeError) as e:
176 logger.warning(f"Error processing infoboxes: {str(e)}")
177 processed_results["infobox"] = None
178
179 # Extract engines used
180 if "engines" in result_data:
181 processed_results["engines"] = result_data["engines"]
182
183 # Extract timing information
184 if "timing" in result_data:
185 processed_results["timing"] = result_data["timing"]
186
187 return processed_results
188
189def detect_factual_intent(message: str) -> bool:
190 """
191 Detect if a message has factual intent (i.e., asking for information).
192
193 Args:
194 message: The message to check
195
196 Returns:
197 bool: True if the message has factual intent, False otherwise
198 """
199 # Simple heuristic: check if the message contains question words or phrases
200 question_words = [
201 "what", "who", "where", "when", "why", "how", "is", "are", "can", "could",
202 "would", "should", "tell me about", "explain", "search", "find", "look up"
203 ]
204
205 message_lower = message.lower()
206
207 # Check for question words
208 for word in question_words:
209 if word in message_lower:
210 return True
211
212 # Check for question marks
213 if "?" in message:
214 return True
215
216 return False
217
218def extract_search_context(results: Dict[str, Any]) -> str:
219 """
220 Extract search context from OpenSearchNG results.
221
222 Args:
223 results: OpenSearchNG results
224
225 Returns:
226 str: Search context
227 """
228 if not results or "results" not in results or not results["results"]:
229 return ""
230
231 # Extract information from results
232 context_parts = []
233
234 # Add answer if available
235 if results.get("answer"):
236 context_parts.append(f"Answer: {results['answer']}")
237
238 # Add infobox if available
239 if results.get("infobox"):
240 infobox = results["infobox"]
241 if isinstance(infobox, dict):
242 infobox_content = infobox.get("content", "")
243 if infobox_content:
244 context_parts.append(f"Infobox: {infobox_content}")
245
246 # Add top results
247 for i, result in enumerate(results["results"][:5], 1):
248 title = result.get("title", "No title")
249 content = result.get("content", "No content")
250 url = result.get("url", "No URL")
251
252 context_parts.append(f"Result {i}: {title}")
253 context_parts.append(f"Content: {content}")
254 context_parts.append(f"URL: {url}")
255 context_parts.append("")
256
257 # Add suggestion if available
258 if results.get("suggestion"):
259 context_parts.append(f"Did you mean: {results['suggestion']}?")
260
261 return "\n".join(context_parts)
262
263def test_direct_opensearch():
264 """Test OpenSearchNG functionality directly."""
265 try:
266 # Test the search command
267 query = "python programming"
268 logger.info(f"Searching for: {query}")
269
270 # Execute the search query
271 results = opensearch_query(query)
272
273 # Check if we got results
274 if not results or not results.get("results"):
275 logger.error(f"No results found for: {query}")
276 return False
277
278 # Print the number of results
279 num_results = len(results.get("results", []))
280 logger.info(f"Got {num_results} results for '{query}'")
281
282 # Print the first result
283 if num_results > 0:
284 first_result = results["results"][0]
285 logger.info(f"First result: {json.dumps(first_result, indent=2)}")
286
287 # Test factual intent detection and search context extraction
288 logger.info("Testing factual intent detection and search context extraction...")
289
290 message = "What is Python programming?"
291 logger.info(f"Message: {message}")
292
293 # Check if the message has factual intent
294 factual_intent = detect_factual_intent(message)
295 if factual_intent:
296 logger.info(f"Factual intent detected in message: {message}")
297
298 # Execute the search query
299 results = opensearch_query(message)
300
301 # Extract search context
302 search_context = extract_search_context(results)
303
304 if search_context:
305 logger.info(f"Search context: {search_context[:100]}...")
306 logger.info("Factual intent detection and search context extraction test completed")
307 else:
308 logger.error("No search context extracted")
309 return False
310 else:
311 logger.error("No factual intent detected")
312 return False
313
314 return True
315
316 except Exception as e:
317 logger.error(f"Error: {e}")
318 return False
319
320if __name__ == "__main__":
321 success = test_direct_opensearch()
322 if success:
323 logger.info("Direct OpenSearchNG test successful!")
324 sys.exit(0)
325 else:
326 logger.error("Direct OpenSearchNG test failed!")
327 sys.exit(1)
ollamatools-new/test_matrix_cmd_search.py +221 -0
……@@ -0,0 +1,221 @@
1#!/usr/bin/env python3
2"""
3Test script to verify search functionality in matrix_cmd.
4"""
5
6import asyncio
7import logging
8import sys
9from matrix_cmd.matrix_client import MatrixClient
10
11# Set up logging
12logging.basicConfig(level=logging.INFO)
13logger = logging.getLogger(__name__)
14
15# Mock class for testing
16class MockMatrixClient:
17 """Mock Matrix client for testing."""
18
19 def __init__(self, enable_search=True, opensearch_config=None):
20 """Initialize the mock Matrix client."""
21 self.enable_search = enable_search
22 self.opensearch_config = opensearch_config or {}
23 self.messages = []
24
25 async def send_message(self, room_id, message):
26 """Send a message to a room."""
27 logger.info(f"Sending message to room {room_id}: {message[:100]}...")
28 self.messages.append({"room_id": room_id, "message": message})
29 return True
30
31# OpenSearchNG configuration
32OPENSEARCH_CONFIG = {
33 "base_url": "http://localhost:8081", # Default local OpenSearchNG URL
34 "search_path": "/search", # Path for search endpoint
35 "preferences_path": "/preferences", # Path for preferences endpoint
36 "categories": ["general"], # Default search categories
37 "engines": ["google", "bing", "duckduckgo"], # Default search engines
38 "language": "en-US", # Default language
39 "format": "json", # Response format (html, json, csv, rss)
40 "timeout": 10, # Default timeout in seconds
41 "num_results": 10, # Default number of results
42 "safe_search": 0 # 0=off, 1=moderate, 2=strict
43}
44
45async def test_matrix_cmd_search():
46 """Test the search functionality in matrix_cmd."""
47 try:
48 # Create a mock Matrix client with search enabled
49 client = MockMatrixClient(
50 enable_search=True,
51 opensearch_config=OPENSEARCH_CONFIG
52 )
53
54 # Test the search command directly
55 logger.info("Testing search command...")
56
57 # Create a mock room ID
58 room_id = "test_room"
59
60 # Test the search command
61 query = "python programming"
62 logger.info(f"Searching for: {query}")
63
64 # Create a mock opensearch_query function to avoid circular imports
65 def mock_opensearch_query(query, config):
66 """Mock implementation of opensearch_query to avoid circular imports."""
67 logger.info(f"Mock OpenSearchNG query: {query}")
68 # Return mock search results
69 return {
70 "query": query,
71 "results": [
72 {
73 "title": "Python Programming - Wikipedia",
74 "url": "https://en.wikipedia.org/wiki/Python_(programming_language)",
75 "content": "Python is a high-level, general-purpose programming language. Its design philosophy emphasizes code readability with the use of significant indentation."
76 },
77 {
78 "title": "Python Tutorial - W3Schools",
79 "url": "https://www.w3schools.com/python/",
80 "content": "Python is a popular programming language. Python can be used on a server to create web applications."
81 },
82 {
83 "title": "Learn Python - Codecademy",
84 "url": "https://www.codecademy.com/learn/learn-python",
85 "content": "Learn Python, a powerful language used by sites like YouTube and Dropbox. Learn the fundamentals of programming to build web apps and manipulate data."
86 }
87 ],
88 "engines": ["mock_engine"],
89 "timing": {"total": 0.1}
90 }
91
92 # Execute the mock search query
93 logger.info(f"Making mock OpenSearchNG query: {query}")
94 results = mock_opensearch_query(query, OPENSEARCH_CONFIG)
95
96 # Format the results
97 formatted_results = ["Search Results:"]
98
99 for i, result in enumerate(results["results"][:5], 1):
100 title = result.get("title", "No title")
101 url = result.get("url", "No URL")
102 content = result.get("content", "No content")
103
104 # Truncate content if it's too long
105 if len(content) > 200:
106 content = content[:200] + "..."
107
108 formatted_results.append(f"{i}. {title}")
109 formatted_results.append(f" URL: {url}")
110 formatted_results.append(f" {content}")
111 formatted_results.append("")
112
113 # Send the formatted results
114 await client.send_message(room_id, "\n".join(formatted_results))
115
116 logger.info("Search command test completed")
117
118 # Test factual intent detection and search context extraction
119 logger.info("Testing factual intent detection and search context extraction...")
120
121 message = "What is Python programming?"
122 logger.info(f"Message: {message}")
123
124 # Define a simple factual intent detection function
125 def detect_factual_intent(message):
126 """Detect if a message has factual intent."""
127 question_words = [
128 "what", "who", "where", "when", "why", "how", "is", "are", "can", "could",
129 "would", "should", "tell me about", "explain", "search", "find", "look up"
130 ]
131
132 message_lower = message.lower()
133
134 # Check for question words
135 for word in question_words:
136 if word in message_lower:
137 return True
138
139 # Check for question marks
140 if "?" in message:
141 return True
142
143 return False
144
145 # Define a simple search context extraction function
146 def extract_search_context(results):
147 """Extract search context from OpenSearchNG results."""
148 if not results or "results" not in results or not results["results"]:
149 return ""
150
151 # Extract information from results
152 context_parts = []
153
154 # Add answer if available
155 if results.get("answer"):
156 context_parts.append(f"Answer: {results['answer']}")
157
158 # Add infobox if available
159 if results.get("infobox"):
160 infobox = results["infobox"]
161 if isinstance(infobox, dict):
162 infobox_content = infobox.get("content", "")
163 if infobox_content:
164 context_parts.append(f"Infobox: {infobox_content}")
165
166 # Add top results
167 for i, result in enumerate(results["results"][:5], 1):
168 title = result.get("title", "No title")
169 content = result.get("content", "No content")
170 url = result.get("url", "No URL")
171
172 context_parts.append(f"Result {i}: {title}")
173 context_parts.append(f"Content: {content}")
174 context_parts.append(f"URL: {url}")
175 context_parts.append("")
176
177 # Add suggestion if available
178 if results.get("suggestion"):
179 context_parts.append(f"Did you mean: {results['suggestion']}?")
180
181 return "\n".join(context_parts)
182
183 # Check if the message has factual intent
184 factual_intent = detect_factual_intent(message)
185 if factual_intent:
186 logger.info(f"Factual intent detected in message: {message}")
187
188 # Execute the mock search query
189 logger.info(f"Making mock OpenSearchNG query: {message}")
190 results = mock_opensearch_query(message, OPENSEARCH_CONFIG)
191
192 # Extract search context
193 search_context = extract_search_context(results)
194
195 if search_context:
196 logger.info(f"Search context: {search_context[:100]}...")
197 logger.info("Factual intent detection and search context extraction test completed")
198 else:
199 logger.error("No search context extracted")
200 return False
201 else:
202 logger.error("No factual intent detected")
203 return False
204
205 return True
206
207 except ImportError as e:
208 logger.error(f"Import error: {e}")
209 return False
210 except Exception as e:
211 logger.error(f"Error: {e}")
212 return False
213
214if __name__ == "__main__":
215 success = asyncio.run(test_matrix_cmd_search())
216 if success:
217 logger.info("Matrix CMD search test successful!")
218 sys.exit(0)
219 else:
220 logger.error("Matrix CMD search test failed!")
221 sys.exit(1)
ollamatools-new/tests/test_matrix_opensearch.py +405 -0
……@@ -0,0 +1,405 @@
1#!/usr/bin/env python3
2"""
3Test script to verify OpenSearchNG functionality in a Matrix client context.
4"""
5
6import logging
7import json
8import sys
9import requests
10from typing import Dict, Any, Optional, List
11
12# Set up logging
13logging.basicConfig(level=logging.INFO)
14logger = logging.getLogger(__name__)
15
16# Default OpenSearchNG configuration
17DEFAULT_OPENSEARCH_CONFIG = {
18 "base_url": "http://localhost:8081", # Default local OpenSearchNG URL
19 "search_path": "/search", # Path for search endpoint
20 "preferences_path": "/preferences", # Path for preferences endpoint
21 "categories": ["general"], # Default search categories
22 "engines": ["google", "bing", "duckduckgo"], # Default search engines
23 "language": "en-US", # Default language
24 "format": "json", # Response format (html, json, csv, rss)
25 "timeout": 10, # Default timeout in seconds
26 "num_results": 10, # Default number of results
27 "safe_search": 0 # 0=off, 1=moderate, 2=strict
28}
29
30def opensearch_query(
31 query: str,
32 config: Optional[Dict[str, Any]] = None,
33) -> Dict[str, Any]:
34 """
35 Perform a search query using a local OpenSearchNG instance.
36
37 Args:
38 query: Search query string
39 config: OpenSearchNG configuration (defaults used if None)
40
41 Returns:
42 dict: Search results with metadata
43 """
44 # Validate input
45 if not query or not isinstance(query, str):
46 raise ValueError("Query must be a non-empty string")
47
48 # Use default config if none provided
49 search_config = DEFAULT_OPENSEARCH_CONFIG.copy()
50 if config:
51 search_config.update(config)
52
53 # Construct the search URL
54 base_url = search_config["base_url"]
55 search_path = search_config["search_path"]
56 search_url = f"{base_url}{search_path}"
57
58 # Prepare query parameters
59 params = {
60 "q": query,
61 "format": search_config["format"],
62 "categories": ",".join(search_config["categories"]),
63 "engines": ",".join(search_config["engines"]),
64 "language": search_config["language"],
65 "safesearch": search_config["safe_search"],
66 "time_range": "", # Optional time range
67 "pageno": 1, # First page of results
68 }
69
70 # Create session with appropriate headers
71 with requests.Session() as session:
72 headers = {
73 "User-Agent": "OllamaTools/0.2.0",
74 "Accept": "application/json"
75 }
76 session.headers.update(headers)
77
78 try:
79 logger.info(f"Making OpenSearchNG query: {query}")
80 response = session.get(
81 search_url,
82 params=params,
83 timeout=search_config["timeout"]
84 )
85
86 # Check for successful response
87 if response.status_code != 200:
88 logger.error(f"OpenSearchNG query failed with status code: {response.status_code}")
89 return {"error": f"Status code: {response.status_code}", "results": []}
90
91 # Parse the response based on format
92 if search_config["format"] == "json":
93 try:
94 result_data = response.json()
95 except json.JSONDecodeError:
96 logger.error("Failed to parse JSON response from OpenSearchNG")
97 return {"error": "JSON decode error", "results": []}
98 else:
99 # For other formats, return the raw content
100 result_data = {"raw_content": response.text}
101
102 # Process the results
103 processed_results = process_opensearch_results(result_data, query)
104
105 return processed_results
106
107 except requests.RequestException as e:
108 logger.error(f"OpenSearchNG request failed: {str(e)}")
109 return {"error": str(e), "results": []}
110
111def process_opensearch_results(
112 result_data: Dict[str, Any],
113 query: str
114) -> Dict[str, Any]:
115 """
116 Process and format the OpenSearchNG results.
117
118 Args:
119 result_data: Raw result data from OpenSearchNG
120 query: Original search query
121
122 Returns:
123 dict: Processed search results with metadata
124 """
125 processed_results = {
126 "query": query,
127 "results": [],
128 "suggestion": None,
129 "answer": None,
130 "infobox": None,
131 "engines": [],
132 "timing": {}
133 }
134
135 # Extract search results
136 if "results" in result_data:
137 for result in result_data["results"]:
138 processed_result = {
139 "title": result.get("title", ""),
140 "url": result.get("url", ""),
141 "content": result.get("content", ""),
142 "engine": result.get("engine", ""),
143 "score": result.get("score", 0),
144 "category": result.get("category", ""),
145 "pretty_url": result.get("pretty_url", "")
146 }
147 processed_results["results"].append(processed_result)
148
149 # Extract suggestions
150 if "suggestions" in result_data and result_data["suggestions"]:
151 try:
152 suggestion = result_data["suggestions"][0]
153 # Handle case where suggestion is a dict instead of a string
154 if isinstance(suggestion, dict):
155 processed_results["suggestion"] = suggestion.get("value", str(suggestion))
156 else:
157 processed_results["suggestion"] = suggestion
158 except (IndexError, TypeError) as e:
159 logger.warning(f"Error processing suggestions: {str(e)}")
160 processed_results["suggestion"] = None
161
162 # Extract answer if available
163 if "answers" in result_data and result_data["answers"]:
164 try:
165 answer = result_data["answers"][0]
166 # Handle case where answer is a dict instead of a string
167 if isinstance(answer, dict):
168 processed_results["answer"] = answer.get("content", str(answer))
169 else:
170 processed_results["answer"] = answer
171 except (IndexError, TypeError) as e:
172 logger.warning(f"Error processing answers: {str(e)}")
173 processed_results["answer"] = None
174
175 # Extract infobox if available
176 if "infoboxes" in result_data and result_data["infoboxes"]:
177 try:
178 infobox = result_data["infoboxes"][0]
179 # Handle case where infobox is a dict (which is expected)
180 processed_results["infobox"] = infobox
181 except (IndexError, TypeError) as e:
182 logger.warning(f"Error processing infoboxes: {str(e)}")
183 processed_results["infobox"] = None
184
185 # Extract engines used
186 if "engines" in result_data:
187 processed_results["engines"] = result_data["engines"]
188
189 # Extract timing information
190 if "timing" in result_data:
191 processed_results["timing"] = result_data["timing"]
192
193 return processed_results
194
195def detect_factual_intent(message: str) -> bool:
196 """
197 Detect if a message has factual intent (i.e., asking for information).
198
199 Args:
200 message: The message to check
201
202 Returns:
203 bool: True if the message has factual intent, False otherwise
204 """
205 # Simple heuristic: check if the message contains question words or phrases
206 question_words = [
207 "what", "who", "where", "when", "why", "how", "is", "are", "can", "could",
208 "would", "should", "tell me about", "explain", "search", "find", "look up"
209 ]
210
211 message_lower = message.lower()
212
213 # Check for question words
214 for word in question_words:
215 if word in message_lower:
216 return True
217
218 # Check for question marks
219 if "?" in message:
220 return True
221
222 return False
223
224def extract_search_context(results: Dict[str, Any]) -> str:
225 """
226 Extract search context from OpenSearchNG results.
227
228 Args:
229 results: OpenSearchNG results
230
231 Returns:
232 str: Search context
233 """
234 if not results or "results" not in results or not results["results"]:
235 return ""
236
237 # Extract information from results
238 context_parts = []
239
240 # Add answer if available
241 if results.get("answer"):
242 context_parts.append(f"Answer: {results['answer']}")
243
244 # Add infobox if available
245 if results.get("infobox"):
246 infobox = results["infobox"]
247 if isinstance(infobox, dict):
248 infobox_content = infobox.get("content", "")
249 if infobox_content:
250 context_parts.append(f"Infobox: {infobox_content}")
251
252 # Add top results
253 for i, result in enumerate(results["results"][:5], 1):
254 title = result.get("title", "No title")
255 content = result.get("content", "No content")
256 url = result.get("url", "No URL")
257
258 context_parts.append(f"Result {i}: {title}")
259 context_parts.append(f"Content: {content}")
260 context_parts.append(f"URL: {url}")
261 context_parts.append("")
262
263 # Add suggestion if available
264 if results.get("suggestion"):
265 context_parts.append(f"Did you mean: {results['suggestion']}?")
266
267 return "\n".join(context_parts)
268
269async def execute_search_for_message(message: str, opensearch_config: Dict[str, Any]) -> str:
270 """
271 Execute a search for a message and return the search context.
272
273 Args:
274 message: The message to search for
275 opensearch_config: The OpenSearchNG configuration
276
277 Returns:
278 str: The search context, or an empty string if search failed
279 """
280 factual_intent = detect_factual_intent(message)
281 if factual_intent:
282 logger.info(f"Factual intent detected in message: {message}")
283 try:
284 results = opensearch_query(message, opensearch_config)
285 search_context = extract_search_context(results)
286 logger.info(f"Search context: {search_context[:100]}...")
287 return search_context
288 except Exception as e:
289 logger.error(f"Error executing search: {e}")
290 return ""
291 else:
292 logger.debug(f"No factual intent detected in message: {message}")
293 return ""
294
295class MockMatrixClient:
296 """Mock Matrix client for testing."""
297
298 def __init__(self, opensearch_config: Dict[str, Any] = None):
299 """Initialize the mock Matrix client."""
300 self.opensearch_config = opensearch_config or DEFAULT_OPENSEARCH_CONFIG
301 self.enable_search = True
302 self.messages = []
303
304 async def send_message(self, room_id: str, message: str) -> None:
305 """
306 Send a message to a room.
307
308 Args:
309 room_id: The ID of the room to send the message to
310 message: The message to send
311 """
312 logger.info(f"Sending message to room {room_id}: {message[:50]}...")
313 self.messages.append({"room_id": room_id, "message": message})
314
315async def handle_search_command(client: MockMatrixClient, query: str, room_id: str) -> None:
316 """
317 Handle the /search command.
318
319 Args:
320 client: The MockMatrixClient instance
321 query: The search query
322 room_id: The ID of the room where the command was sent
323 """
324 if not client.enable_search:
325 await client.send_message(room_id, "Search functionality is not enabled.")
326 return
327
328 await client.send_message(room_id, f"Searching for: {query}...")
329
330 try:
331 results = opensearch_query(query, client.opensearch_config)
332
333 if not results or not results.get("results"):
334 await client.send_message(room_id, f"No results found for: {query}")
335 return
336
337 # Format the results
338 formatted_results = ["Search Results:"]
339
340 for i, result in enumerate(results["results"][:5], 1):
341 title = result.get("title", "No title")
342 url = result.get("url", "No URL")
343 content = result.get("content", "No content")
344
345 # Truncate content if it's too long
346 if len(content) > 200:
347 content = content[:200] + "..."
348
349 formatted_results.append(f"{i}. {title}")
350 formatted_results.append(f" URL: {url}")
351 formatted_results.append(f" {content}")
352 formatted_results.append("")
353
354 # Send the formatted results
355 await client.send_message(room_id, "\n".join(formatted_results))
356
357 except Exception as e:
358 logger.error(f"Error executing search: {e}")
359 await client.send_message(room_id, f"Error executing search: {str(e)}")
360
361async def test_matrix_search():
362 """Test the Matrix search functionality."""
363 try:
364 # Create a mock Matrix client
365 client = MockMatrixClient()
366
367 # Test the search command
368 logger.info("Testing search command...")
369 await handle_search_command(client, "python programming", "test_room")
370
371 # Check if we got a response
372 if not client.messages:
373 logger.error("No messages sent by the client")
374 return False
375
376 # Print the last message
377 last_message = client.messages[-1]["message"]
378 logger.info(f"Last message: {last_message[:100]}...")
379
380 # Test search for message
381 logger.info("Testing search for message...")
382 search_context = await execute_search_for_message("What is Python programming?", client.opensearch_config)
383
384 if not search_context:
385 logger.error("No search context returned")
386 return False
387
388 logger.info(f"Search context: {search_context[:100]}...")
389
390 return True
391
392 except Exception as e:
393 logger.error(f"Error: {e}")
394 return False
395
396if __name__ == "__main__":
397 import asyncio
398
399 success = asyncio.run(test_matrix_search())
400 if success:
401 logger.info("Matrix search test successful!")
402 sys.exit(0)
403 else:
404 logger.error("Matrix search test failed!")
405 sys.exit(1)
ollamatools-new/tests/test_ollamatools_opensearch.py +59 -0
……@@ -0,0 +1,59 @@
1#!/usr/bin/env python3
2"""
3Test script to verify OpenSearchNG integration with OllamaTools.
4"""
5
6import logging
7import json
8import sys
9
10# Set up logging
11logging.basicConfig(level=logging.INFO)
12logger = logging.getLogger(__name__)
13
14def test_opensearch_integration():
15 """Test the OpenSearchNG integration with OllamaTools."""
16 try:
17 # Try to import the OpenSearchNG tools
18 from ollamatools.plugins.integrations.opensearch.tools import (
19 opensearch_query,
20 get_opensearch_info,
21 DEFAULT_OPENSEARCH_CONFIG
22 )
23
24 logger.info("Successfully imported OpenSearchNG tools")
25
26 # Test getting OpenSearchNG info
27 logger.info("Testing OpenSearchNG info...")
28 info_result = get_opensearch_info()
29 logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
30
31 # Test OpenSearchNG query
32 logger.info("Testing OpenSearchNG query...")
33 query_result = opensearch_query("python programming")
34
35 num_results = len(query_result.get("results", []))
36 logger.info(f"Got {num_results} results for 'python programming'")
37
38 # Print the first result
39 if num_results > 0:
40 first_result = query_result["results"][0]
41 logger.info(f"First result: {json.dumps(first_result, indent=2)}")
42
43 return True
44
45 except ImportError as e:
46 logger.error(f"Import error: {e}")
47 return False
48 except Exception as e:
49 logger.error(f"Error: {e}")
50 return False
51
52if __name__ == "__main__":
53 success = test_opensearch_integration()
54 if success:
55 logger.info("OpenSearchNG integration test successful!")
56 sys.exit(0)
57 else:
58 logger.error("OpenSearchNG integration test failed!")
59 sys.exit(1)
ollamatools-new/tests/test_ollamatools_opensearch_direct.py +59 -0
……@@ -0,0 +1,59 @@
1#!/usr/bin/env python3
2"""
3Test script to directly verify OpenSearchNG integration with OllamaTools.
4"""
5
6import logging
7import json
8import sys
9
10# Set up logging
11logging.basicConfig(level=logging.INFO)
12logger = logging.getLogger(__name__)
13
14def test_opensearch_integration():
15 """Test the OpenSearchNG integration with OllamaTools."""
16 try:
17 # Try to import the OpenSearchNG tools directly
18 from ollamatools.plugins.integrations.opensearch.tools import (
19 opensearch_query,
20 get_opensearch_info,
21 DEFAULT_OPENSEARCH_CONFIG
22 )
23
24 logger.info("Successfully imported OpenSearchNG tools")
25
26 # Test getting OpenSearchNG info
27 logger.info("Testing OpenSearchNG info...")
28 info_result = get_opensearch_info()
29 logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
30
31 # Test OpenSearchNG query
32 logger.info("Testing OpenSearchNG query...")
33 query_result = opensearch_query("python programming")
34
35 num_results = len(query_result.get("results", []))
36 logger.info(f"Got {num_results} results for 'python programming'")
37
38 # Print the first result
39 if num_results > 0:
40 first_result = query_result["results"][0]
41 logger.info(f"First result: {json.dumps(first_result, indent=2)}")
42
43 return True
44
45 except ImportError as e:
46 logger.error(f"Import error: {e}")
47 return False
48 except Exception as e:
49 logger.error(f"Error: {e}")
50 return False
51
52if __name__ == "__main__":
53 success = test_opensearch_integration()
54 if success:
55 logger.info("OpenSearchNG integration test successful!")
56 sys.exit(0)
57 else:
58 logger.error("OpenSearchNG integration test failed!")
59 sys.exit(1)
ollamatools-new/tests/test_opensearch.py +146 -0
……@@ -0,0 +1,146 @@
1#!/usr/bin/env python3
2"""
3Test script to verify OpenSearchNG functionality.
4"""
5
6import requests
7import json
8import logging
9from typing import Dict, Any
10
11# Set up logging
12logging.basicConfig(level=logging.INFO)
13logger = logging.getLogger(__name__)
14
15def test_opensearch_query(query: str = "test") -> Dict[str, Any]:
16 """
17 Test OpenSearchNG query functionality.
18
19 Args:
20 query: The search query
21
22 Returns:
23 Dict[str, Any]: Search results
24 """
25 # Default OpenSearchNG configuration
26 search_config = {
27 "base_url": "http://localhost:8081", # Default local OpenSearchNG URL
28 "search_path": "/search", # Path for search endpoint
29 "format": "json", # Response format
30 "categories": ["general"], # Default search categories
31 "engines": ["google", "bing", "duckduckgo"], # Default search engines
32 "language": "en-US", # Default language
33 "timeout": 10, # Default timeout in seconds
34 }
35
36 # Construct the search URL
37 base_url = search_config["base_url"]
38 search_path = search_config["search_path"]
39 search_url = f"{base_url}{search_path}"
40
41 # Prepare query parameters
42 params = {
43 "q": query,
44 "format": search_config["format"],
45 "categories": ",".join(search_config["categories"]),
46 "engines": ",".join(search_config["engines"]),
47 "language": search_config["language"],
48 }
49
50 # Create session with appropriate headers
51 with requests.Session() as session:
52 headers = {
53 "User-Agent": "OllamaTools/0.2.0",
54 "Accept": "application/json"
55 }
56 session.headers.update(headers)
57
58 try:
59 logger.info(f"Making OpenSearchNG query: {query}")
60 response = session.get(
61 search_url,
62 params=params,
63 timeout=search_config["timeout"]
64 )
65
66 # Check for successful response
67 if response.status_code != 200:
68 logger.error(f"OpenSearchNG query failed with status code: {response.status_code}")
69 return {"error": f"Status code: {response.status_code}", "results": []}
70
71 # Parse the response
72 try:
73 result_data = response.json()
74 logger.info(f"Got {len(result_data.get('results', []))} results")
75 return result_data
76 except json.JSONDecodeError:
77 logger.error("Failed to parse JSON response from OpenSearchNG")
78 return {"error": "JSON decode error", "results": []}
79
80 except requests.RequestException as e:
81 logger.error(f"OpenSearchNG request failed: {str(e)}")
82 return {"error": str(e), "results": []}
83
84def test_opensearch_info() -> Dict[str, Any]:
85 """
86 Test getting information about the OpenSearchNG instance.
87
88 Returns:
89 Dict[str, Any]: Information about the OpenSearchNG instance
90 """
91 # Default configuration
92 search_config = {
93 "base_url": "http://localhost:8081",
94 "timeout": 10,
95 }
96
97 # Construct the URL
98 base_url = search_config["base_url"]
99 config_url = f"{base_url}/config"
100
101 try:
102 logger.info("Getting OpenSearchNG instance information")
103 response = requests.get(
104 config_url,
105 timeout=search_config["timeout"]
106 )
107
108 # Check for successful response
109 if response.status_code != 200:
110 logger.error(f"OpenSearchNG info request failed with status code: {response.status_code}")
111 return {"error": f"Status code: {response.status_code}"}
112
113 # Parse the response
114 try:
115 config_data = response.json()
116 return config_data
117 except json.JSONDecodeError:
118 logger.error("Failed to parse JSON response from OpenSearchNG")
119 return {"error": "JSON decode error"}
120
121 except requests.RequestException as e:
122 logger.error(f"OpenSearchNG info request failed: {str(e)}")
123 return {"error": str(e)}
124
125if __name__ == "__main__":
126 # Test OpenSearchNG info
127 logger.info("Testing OpenSearchNG info...")
128 info_result = test_opensearch_info()
129 if "error" in info_result:
130 logger.error(f"Error getting OpenSearchNG info: {info_result['error']}")
131 else:
132 logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
133
134 # Test OpenSearchNG query
135 logger.info("Testing OpenSearchNG query...")
136 query_result = test_opensearch_query("python programming")
137 if "error" in query_result:
138 logger.error(f"Error executing query: {query_result['error']}")
139 else:
140 num_results = len(query_result.get("results", []))
141 logger.info(f"Got {num_results} results for 'python programming'")
142
143 # Print the first result
144 if num_results > 0:
145 first_result = query_result["results"][0]
146 logger.info(f"First result: {json.dumps(first_result, indent=2)}")
ollamatools-new/tests/test_opensearch_direct.py +282 -0
……@@ -0,0 +1,282 @@
1#!/usr/bin/env python3
2"""
3Test script to directly verify OpenSearchNG functionality without importing from ollamatools.
4"""
5
6import logging
7import json
8import sys
9import requests
10from typing import Dict, Any, Optional, Set
11
12# Set up logging
13logging.basicConfig(level=logging.INFO)
14logger = logging.getLogger(__name__)
15
16# Default OpenSearchNG configuration
17DEFAULT_OPENSEARCH_CONFIG = {
18 "base_url": "http://localhost:8081", # Default local OpenSearchNG URL
19 "search_path": "/search", # Path for search endpoint
20 "preferences_path": "/preferences", # Path for preferences endpoint
21 "categories": ["general"], # Default search categories
22 "engines": ["google", "bing", "duckduckgo"], # Default search engines
23 "language": "en-US", # Default language
24 "format": "json", # Response format (html, json, csv, rss)
25 "timeout": 10, # Default timeout in seconds
26 "num_results": 10, # Default number of results
27 "safe_search": 0 # 0=off, 1=moderate, 2=strict
28}
29
30def opensearch_query(
31 query: str,
32 config: Optional[Dict[str, Any]] = None,
33) -> Dict[str, Any]:
34 """
35 Perform a search query using a local OpenSearchNG instance.
36
37 Args:
38 query: Search query string
39 config: OpenSearchNG configuration (defaults used if None)
40
41 Returns:
42 dict: Search results with metadata
43 """
44 # Validate input
45 if not query or not isinstance(query, str):
46 raise ValueError("Query must be a non-empty string")
47
48 # Use default config if none provided
49 search_config = DEFAULT_OPENSEARCH_CONFIG.copy()
50 if config:
51 search_config.update(config)
52
53 # Construct the search URL
54 base_url = search_config["base_url"]
55 search_path = search_config["search_path"]
56 search_url = f"{base_url}{search_path}"
57
58 # Prepare query parameters
59 params = {
60 "q": query,
61 "format": search_config["format"],
62 "categories": ",".join(search_config["categories"]),
63 "engines": ",".join(search_config["engines"]),
64 "language": search_config["language"],
65 "safesearch": search_config["safe_search"],
66 "time_range": "", # Optional time range
67 "pageno": 1, # First page of results
68 }
69
70 # Create session with appropriate headers
71 with requests.Session() as session:
72 headers = {
73 "User-Agent": "OllamaTools/0.2.0",
74 "Accept": "application/json"
75 }
76 session.headers.update(headers)
77
78 try:
79 logger.info(f"Making OpenSearchNG query: {query}")
80 response = session.get(
81 search_url,
82 params=params,
83 timeout=search_config["timeout"]
84 )
85
86 # Check for successful response
87 if response.status_code != 200:
88 logger.error(f"OpenSearchNG query failed with status code: {response.status_code}")
89 return {"error": f"Status code: {response.status_code}", "results": []}
90
91 # Parse the response based on format
92 if search_config["format"] == "json":
93 try:
94 result_data = response.json()
95 except json.JSONDecodeError:
96 logger.error("Failed to parse JSON response from OpenSearchNG")
97 return {"error": "JSON decode error", "results": []}
98 else:
99 # For other formats, return the raw content
100 result_data = {"raw_content": response.text}
101
102 # Process the results
103 processed_results = process_opensearch_results(result_data, query)
104
105 return processed_results
106
107 except requests.RequestException as e:
108 logger.error(f"OpenSearchNG request failed: {str(e)}")
109 return {"error": str(e), "results": []}
110
111def process_opensearch_results(
112 result_data: Dict[str, Any],
113 query: str
114) -> Dict[str, Any]:
115 """
116 Process and format the OpenSearchNG results.
117
118 Args:
119 result_data: Raw result data from OpenSearchNG
120 query: Original search query
121
122 Returns:
123 dict: Processed search results with metadata
124 """
125 processed_results = {
126 "query": query,
127 "results": [],
128 "suggestion": None,
129 "answer": None,
130 "infobox": None,
131 "engines": [],
132 "timing": {}
133 }
134
135 # Extract search results
136 if "results" in result_data:
137 for result in result_data["results"]:
138 processed_result = {
139 "title": result.get("title", ""),
140 "url": result.get("url", ""),
141 "content": result.get("content", ""),
142 "engine": result.get("engine", ""),
143 "score": result.get("score", 0),
144 "category": result.get("category", ""),
145 "pretty_url": result.get("pretty_url", "")
146 }
147 processed_results["results"].append(processed_result)
148
149 # Extract suggestions
150 if "suggestions" in result_data and result_data["suggestions"]:
151 try:
152 suggestion = result_data["suggestions"][0]
153 # Handle case where suggestion is a dict instead of a string
154 if isinstance(suggestion, dict):
155 processed_results["suggestion"] = suggestion.get("value", str(suggestion))
156 else:
157 processed_results["suggestion"] = suggestion
158 except (IndexError, TypeError) as e:
159 logger.warning(f"Error processing suggestions: {str(e)}")
160 processed_results["suggestion"] = None
161
162 # Extract answer if available
163 if "answers" in result_data and result_data["answers"]:
164 try:
165 answer = result_data["answers"][0]
166 # Handle case where answer is a dict instead of a string
167 if isinstance(answer, dict):
168 processed_results["answer"] = answer.get("content", str(answer))
169 else:
170 processed_results["answer"] = answer
171 except (IndexError, TypeError) as e:
172 logger.warning(f"Error processing answers: {str(e)}")
173 processed_results["answer"] = None
174
175 # Extract infobox if available
176 if "infoboxes" in result_data and result_data["infoboxes"]:
177 try:
178 infobox = result_data["infoboxes"][0]
179 # Handle case where infobox is a dict (which is expected)
180 processed_results["infobox"] = infobox
181 except (IndexError, TypeError) as e:
182 logger.warning(f"Error processing infoboxes: {str(e)}")
183 processed_results["infobox"] = None
184
185 # Extract engines used
186 if "engines" in result_data:
187 processed_results["engines"] = result_data["engines"]
188
189 # Extract timing information
190 if "timing" in result_data:
191 processed_results["timing"] = result_data["timing"]
192
193 return processed_results
194
195def get_opensearch_info(
196 config: Optional[Dict[str, Any]] = None,
197) -> Dict[str, Any]:
198 """
199 Get information about the OpenSearchNG instance.
200
201 Args:
202 config: OpenSearchNG configuration (defaults used if None)
203
204 Returns:
205 dict: Information about the OpenSearchNG instance
206 """
207 # Use default config if none provided
208 search_config = DEFAULT_OPENSEARCH_CONFIG.copy()
209 if config:
210 search_config.update(config)
211
212 # Construct the URL
213 base_url = search_config["base_url"]
214 config_url = f"{base_url}/config"
215
216 try:
217 logger.info("Getting OpenSearchNG instance information")
218 response = requests.get(
219 config_url,
220 timeout=search_config["timeout"]
221 )
222
223 # Check for successful response
224 if response.status_code != 200:
225 logger.error(f"OpenSearchNG info request failed with status code: {response.status_code}")
226 return {"error": f"Status code: {response.status_code}"}
227
228 # Parse the response
229 try:
230 config_data = response.json()
231 return config_data
232 except json.JSONDecodeError:
233 logger.error("Failed to parse JSON response from OpenSearchNG")
234 return {"error": "JSON decode error"}
235
236 except requests.RequestException as e:
237 logger.error(f"OpenSearchNG info request failed: {str(e)}")
238 return {"error": str(e)}
239
240def test_opensearch_functionality():
241 """Test the OpenSearchNG functionality."""
242 try:
243 # Test getting OpenSearchNG info
244 logger.info("Testing OpenSearchNG info...")
245 info_result = get_opensearch_info()
246
247 if "error" in info_result:
248 logger.error(f"Error getting OpenSearchNG info: {info_result['error']}")
249 return False
250 else:
251 logger.info(f"OpenSearchNG info: {json.dumps(info_result, indent=2)[:200]}...")
252
253 # Test OpenSearchNG query
254 logger.info("Testing OpenSearchNG query...")
255 query_result = opensearch_query("python programming")
256
257 if "error" in query_result:
258 logger.error(f"Error executing query: {query_result['error']}")
259 return False
260
261 num_results = len(query_result.get("results", []))
262 logger.info(f"Got {num_results} results for 'python programming'")
263
264 # Print the first result
265 if num_results > 0:
266 first_result = query_result["results"][0]
267 logger.info(f"First result: {json.dumps(first_result, indent=2)}")
268
269 return True
270
271 except Exception as e:
272 logger.error(f"Error: {e}")
273 return False
274
275if __name__ == "__main__":
276 success = test_opensearch_functionality()
277 if success:
278 logger.info("OpenSearchNG functionality test successful!")
279 sys.exit(0)
280 else:
281 logger.error("OpenSearchNG functionality test failed!")
282 sys.exit(1)