feat: add es headers config (#3088)

mem0/configs/vector_stores/elasticsearch.py

@@ -19,6 +19,7 @@ class ElasticsearchConfig(BaseModel):
     custom_search_query: Optional[Callable[[List[float], int, Optional[Dict]], Dict]] = Field(
         None, description="Custom search query function. Parameters: (query, limit, filters) -> Dict"
     )
+    headers: Optional[Dict[str, str]] = Field(None, description="Custom headers to include in requests")
 
     @model_validator(mode="before")
     @classmethod
@@ -33,6 +34,23 @@ class ElasticsearchConfig(BaseModel):
 
         return values
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_headers(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate headers format and content"""
+        headers = values.get("headers")
+        if headers is not None:
+            # Check if headers is a dictionary
+            if not isinstance(headers, dict):
+                raise ValueError("headers must be a dictionary")
+            
+            # Check if all keys and values are strings
+            for key, value in headers.items():
+                if not isinstance(key, str) or not isinstance(value, str):
+                    raise ValueError("All header keys and values must be strings")
+        
+        return values
+
     @model_validator(mode="before")
     @classmethod
     def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:

mem0/vector_stores/elasticsearch.py

@@ -31,12 +31,14 @@ class ElasticsearchDB(VectorStoreBase):
                 cloud_id=config.cloud_id,
                 api_key=config.api_key,
                 verify_certs=config.verify_certs,
+                headers= config.headers or {},
             )
         else:
             self.client = Elasticsearch(
                 hosts=[f"{config.host}" if config.port is None else f"{config.host}:{config.port}"],
                 basic_auth=(config.user, config.password) if (config.user and config.password) else None,
                 verify_certs=config.verify_certs,
+                headers= config.headers or {},
             )
 
         self.collection_name = config.collection_name