local-deep-research 0.4.4__py3-none-any.whl → 0.5.2__py3-none-any.whl

local_deep_research/utilities/es_utils.py

@@ -2,10 +2,9 @@
 Elasticsearch utilities for indexing and managing documents.
 """
 
-import json
 import logging
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional
 
 from elasticsearch import Elasticsearch
 from elasticsearch.helpers import bulk
@@ -16,11 +15,11 @@ logger = logging.getLogger(__name__)
 class ElasticsearchManager:
     """
     Utility class for managing Elasticsearch indices and documents.
-    
+
     This class provides methods for creating indices, indexing documents,
     and performing other Elasticsearch management tasks.
     """
-    
+
     def __init__(
         self,
         hosts: List[str] = ["http://localhost:9200"],
@@ -31,7 +30,7 @@ class ElasticsearchManager:
     ):
         """
         Initialize the Elasticsearch manager.
-        
+
         Args:
             hosts: List of Elasticsearch hosts
             username: Optional username for authentication
@@ -41,31 +40,37 @@ class ElasticsearchManager:
         """
         # Initialize the Elasticsearch client
         es_args = {}
-        
+
         # Basic authentication
         if username and password:
             es_args["basic_auth"] = (username, password)
-        
+
         # API key authentication
         if api_key:
             es_args["api_key"] = api_key
-        
+
         # Cloud ID for Elastic Cloud
         if cloud_id:
             es_args["cloud_id"] = cloud_id
-        
+
         # Connect to Elasticsearch
         self.client = Elasticsearch(hosts, **es_args)
-        
+
         # Verify connection
         try:
             info = self.client.info()
-            logger.info(f"Connected to Elasticsearch cluster: {info.get('cluster_name')}")
-            logger.info(f"Elasticsearch version: {info.get('version', {}).get('number')}")
+            logger.info(
+                f"Connected to Elasticsearch cluster: {info.get('cluster_name')}"
+            )
+            logger.info(
+                f"Elasticsearch version: {info.get('version', {}).get('number')}"
+            )
         except Exception as e:
             logger.error(f"Failed to connect to Elasticsearch: {str(e)}")
-            raise ConnectionError(f"Could not connect to Elasticsearch: {str(e)}")
-    
+            raise ConnectionError(
+                f"Could not connect to Elasticsearch: {str(e)}"
+            )
+
     def create_index(
         self,
         index_name: str,
@@ -74,21 +79,23 @@ class ElasticsearchManager:
     ) -> bool:
         """
         Create an Elasticsearch index with optional mappings and settings.
-        
+
         Args:
             index_name: Name of the index to create
             mappings: Optional mappings for the index fields
             settings: Optional settings for the index
-            
+
         Returns:
             bool: True if successful, False otherwise
         """
         try:
             # Check if index already exists
             if self.client.indices.exists(index=index_name):
-                logger.warning(f"Index '{index_name}' already exists - skipping creation")
+                logger.warning(
+                    f"Index '{index_name}' already exists - skipping creation"
+                )
                 return True
-            
+
             # Default mappings for better text search if none provided
             if mappings is None:
                 mappings = {
@@ -99,83 +106,69 @@ class ElasticsearchManager:
                             "fields": {
                                 "keyword": {
                                     "type": "keyword",
-                                    "ignore_above": 256
+                                    "ignore_above": 256,
                                 }
-                            }
-                        },
-                        "content": {
-                            "type": "text",
-                            "analyzer": "standard"
-                        },
-                        "url": {
-                            "type": "keyword"
-                        },
-                        "source": {
-                            "type": "keyword"
-                        },
-                        "timestamp": {
-                            "type": "date"
+                            },
                         },
-                        "metadata": {
-                            "type": "object",
-                            "enabled": True
-                        }
+                        "content": {"type": "text", "analyzer": "standard"},
+                        "url": {"type": "keyword"},
+                        "source": {"type": "keyword"},
+                        "timestamp": {"type": "date"},
+                        "metadata": {"type": "object", "enabled": True},
                     }
                 }
-            
+
             # Default settings if none provided
             if settings is None:
                 settings = {
                     "number_of_shards": 1,
                     "number_of_replicas": 0,
                     "analysis": {
-                        "analyzer": {
-                            "standard": {
-                                "type": "standard"
-                            }
-                        }
-                    }
+                        "analyzer": {"standard": {"type": "standard"}}
+                    },
                 }
-            
+
             # Create the index with mappings and settings
             create_response = self.client.indices.create(
                 index=index_name,
                 mappings=mappings,
                 settings=settings,
             )
-            
+
             logger.info(f"Created index '{index_name}': {create_response}")
             return True
-            
+
         except Exception as e:
             logger.error(f"Error creating index '{index_name}': {str(e)}")
             return False
-    
+
     def delete_index(self, index_name: str) -> bool:
         """
         Delete an Elasticsearch index.
-        
+
         Args:
             index_name: Name of the index to delete
-            
+
         Returns:
             bool: True if successful, False otherwise
         """
         try:
             # Check if index exists
             if not self.client.indices.exists(index=index_name):
-                logger.warning(f"Index '{index_name}' does not exist - skipping deletion")
+                logger.warning(
+                    f"Index '{index_name}' does not exist - skipping deletion"
+                )
                 return True
-            
+
             # Delete the index
             delete_response = self.client.indices.delete(index=index_name)
             logger.info(f"Deleted index '{index_name}': {delete_response}")
             return True
-            
+
         except Exception as e:
             logger.error(f"Error deleting index '{index_name}': {str(e)}")
             return False
-    
+
     def index_document(
         self,
         index_name: str,
@@ -185,13 +178,13 @@ class ElasticsearchManager:
     ) -> Optional[str]:
         """
         Index a single document in Elasticsearch.
-        
+
         Args:
             index_name: Name of the index to add the document to
             document: The document to index
             document_id: Optional document ID (will be generated if not provided)
             refresh: Whether to refresh the index after indexing
-            
+
         Returns:
             str: Document ID if successful, None otherwise
         """
@@ -203,14 +196,16 @@ class ElasticsearchManager:
                 id=document_id,
                 refresh=refresh,
             )
-            
-            logger.info(f"Indexed document in '{index_name}' with ID: {response['_id']}")
+
+            logger.info(
+                f"Indexed document in '{index_name}' with ID: {response['_id']}"
+            )
             return response["_id"]
-            
+
         except Exception as e:
             logger.error(f"Error indexing document in '{index_name}': {str(e)}")
             return None
-    
+
     def bulk_index_documents(
         self,
         index_name: str,
@@ -220,13 +215,13 @@ class ElasticsearchManager:
     ) -> int:
         """
         Bulk index multiple documents in Elasticsearch.
-        
+
         Args:
             index_name: Name of the index to add the documents to
             documents: List of documents to index
             id_field: Optional field in the documents to use as the document ID
             refresh: Whether to refresh the index after indexing
-            
+
         Returns:
             int: Number of successfully indexed documents
         """
@@ -238,13 +233,13 @@ class ElasticsearchManager:
                     "_index": index_name,
                     "_source": doc,
                 }
-                
+
                 # Use the specified field as the document ID if provided
                 if id_field and id_field in doc:
                     action["_id"] = doc[id_field]
-                
+
                 actions.append(action)
-            
+
             # Execute the bulk indexing
             success, failed = bulk(
                 self.client,
@@ -252,14 +247,18 @@ class ElasticsearchManager:
                 refresh=refresh,
                 stats_only=True,
             )
-            
-            logger.info(f"Bulk indexed {success} documents in '{index_name}', failed: {failed}")
+
+            logger.info(
+                f"Bulk indexed {success} documents in '{index_name}', failed: {failed}"
+            )
             return success
-            
+
         except Exception as e:
-            logger.error(f"Error bulk indexing documents in '{index_name}': {str(e)}")
+            logger.error(
+                f"Error bulk indexing documents in '{index_name}': {str(e)}"
+            )
             return 0
-    
+
     def index_file(
         self,
         index_name: str,
@@ -271,7 +270,7 @@ class ElasticsearchManager:
     ) -> Optional[str]:
         """
         Index a file in Elasticsearch, extracting text content and metadata.
-        
+
         Args:
             index_name: Name of the index to add the document to
             file_path: Path to the file to index
@@ -279,53 +278,59 @@ class ElasticsearchManager:
             title_field: Field name to store the file title (filename if not specified)
             extract_metadata: Whether to extract file metadata
             refresh: Whether to refresh the index after indexing
-            
+
         Returns:
             str: Document ID if successful, None otherwise
         """
         try:
-            from langchain_community.document_loaders import UnstructuredFileLoader
-            
+            from langchain_community.document_loaders import (
+                UnstructuredFileLoader,
+            )
+
             # Extract file content and metadata
             loader = UnstructuredFileLoader(file_path)
             documents = loader.load()
-            
+
             # Combine all content from the documents
             content = "\n\n".join([doc.page_content for doc in documents])
-            
+
             # Get the filename for the title
             filename = os.path.basename(file_path)
             title = filename
-            
+
             # Prepare the document
             document = {
                 content_field: content,
             }
-            
+
             # Add title if requested
             if title_field:
                 document[title_field] = title
-            
+
             # Add metadata if requested
             if extract_metadata and documents:
                 # Include metadata from the first document
                 document["metadata"] = documents[0].metadata
-                
+
                 # Add file-specific metadata
                 document["source"] = file_path
-                document["file_extension"] = os.path.splitext(filename)[1].lstrip(".")
+                document["file_extension"] = os.path.splitext(filename)[
+                    1
+                ].lstrip(".")
                 document["filename"] = filename
-            
+
             # Index the document
             return self.index_document(index_name, document, refresh=refresh)
-            
+
         except ImportError:
-            logger.error("UnstructuredFileLoader not available. Please install the 'unstructured' package.")
+            logger.error(
+                "UnstructuredFileLoader not available. Please install the 'unstructured' package."
+            )
             return None
         except Exception as e:
             logger.error(f"Error indexing file '{file_path}': {str(e)}")
             return None
-    
+
     def index_directory(
         self,
         index_name: str,
@@ -338,7 +343,7 @@ class ElasticsearchManager:
     ) -> int:
         """
         Index all matching files in a directory in Elasticsearch.
-        
+
         Args:
             index_name: Name of the index to add the documents to
             directory_path: Path to the directory containing files to index
@@ -347,22 +352,24 @@ class ElasticsearchManager:
             title_field: Field name to store the file title
             extract_metadata: Whether to extract file metadata
             refresh: Whether to refresh the index after indexing
-            
+
         Returns:
             int: Number of successfully indexed files
         """
         try:
             import glob
-            
+
             # Find all matching files
             all_files = []
             for pattern in file_patterns:
                 pattern_path = os.path.join(directory_path, pattern)
                 matching_files = glob.glob(pattern_path)
                 all_files.extend(matching_files)
-            
-            logger.info(f"Found {len(all_files)} files matching patterns {file_patterns} in {directory_path}")
-            
+
+            logger.info(
+                f"Found {len(all_files)} files matching patterns {file_patterns} in {directory_path}"
+            )
+
             # Index each file
             successful_count = 0
             for file_path in all_files:
@@ -375,17 +382,21 @@ class ElasticsearchManager:
                     extract_metadata=extract_metadata,
                     refresh=refresh,
                 )
-                
+
                 if doc_id:
                     successful_count += 1
-            
-            logger.info(f"Successfully indexed {successful_count} files out of {len(all_files)}")
+
+            logger.info(
+                f"Successfully indexed {successful_count} files out of {len(all_files)}"
+            )
             return successful_count
-            
+
         except Exception as e:
-            logger.error(f"Error indexing directory '{directory_path}': {str(e)}")
+            logger.error(
+                f"Error indexing directory '{directory_path}': {str(e)}"
+            )
             return 0
-    
+
     def search(
         self,
         index_name: str,
@@ -396,14 +407,14 @@ class ElasticsearchManager:
     ) -> Dict[str, Any]:
         """
         Search for documents in Elasticsearch.
-        
+
         Args:
             index_name: Name of the index to search
             query: Search query
             fields: Fields to search in
             size: Maximum number of results to return
             highlight: Whether to include highlighted excerpts in results
-            
+
         Returns:
             Dict: Elasticsearch search response
         """
@@ -419,7 +430,7 @@ class ElasticsearchManager:
                 },
                 "size": size,
             }
-            
+
             # Add highlighting if requested
             if highlight:
                 search_query["highlight"] = {
@@ -427,15 +438,15 @@ class ElasticsearchManager:
                     "pre_tags": ["<em>"],
                     "post_tags": ["</em>"],
                 }
-            
+
             # Execute the search
             response = self.client.search(
                 index=index_name,
                 body=search_query,
             )
-            
+
             return response
-            
+
         except Exception as e:
             logger.error(f"Error searching index '{index_name}': {str(e)}")
-            return {"error": str(e)} 
+            return {"error": str(e)}

local_deep_research/utilities/llm_utils.py

@@ -78,7 +78,9 @@ def get_model(
             api_key = os.getenv("OPENAI_API_KEY")
             if not api_key:
                 raise ValueError("OPENAI_API_KEY environment variable not set")
-            return ChatOpenAI(model=model_name, api_key=api_key, **common_params)
+            return ChatOpenAI(
+                model=model_name, api_key=api_key, **common_params
+            )
         except ImportError:
             logger.error("langchain_openai not available")
             raise
@@ -89,7 +91,9 @@ def get_model(
 
             api_key = os.getenv("ANTHROPIC_API_KEY")
             if not api_key:
-                raise ValueError("ANTHROPIC_API_KEY environment variable not set")
+                raise ValueError(
+                    "ANTHROPIC_API_KEY environment variable not set"
+                )
             return ChatAnthropic(
                 model=model_name, anthropic_api_key=api_key, **common_params
             )
@@ -103,7 +107,9 @@ def get_model(
 
             api_key = os.getenv("OPENAI_ENDPOINT_API_KEY")
             if not api_key:
-                raise ValueError("OPENAI_ENDPOINT_API_KEY environment variable not set")
+                raise ValueError(
+                    "OPENAI_ENDPOINT_API_KEY environment variable not set"
+                )
 
             endpoint_url = kwargs.get(
                 "OPENAI_ENDPOINT_URL", "https://openrouter.ai/api/v1"
@@ -113,7 +119,9 @@ def get_model(
                 "OPENAI_ENDPOINT_REQUIRES_MODEL", True
             ):
                 return ChatOpenAI(
-                    api_key=api_key, openai_api_base=endpoint_url, **common_params
+                    api_key=api_key,
+                    openai_api_base=endpoint_url,
+                    **common_params,
                 )
             else:
                 return ChatOpenAI(
@@ -130,7 +138,9 @@ def get_model(
     try:
         from langchain_ollama import ChatOllama
 
-        logger.warning(f"Unknown model type '{model_type}', defaulting to Ollama")
+        logger.warning(
+            f"Unknown model type '{model_type}', defaulting to Ollama"
+        )
         return ChatOllama(model=model_name, **common_params)
     except (ImportError, Exception) as e:
         logger.error(f"Failed to load any model: {e}")