Merge pull request #62 from planetf1/customheader

Author: planetf1

README.md

@@ -442,6 +442,16 @@ maestro-knowledge/
 - `OPENAI_API_KEY`: Required for OpenAI embedding models
 - `MAESTRO_KNOWLEDGE_MCP_SERVER_URI`: MCP server URI for CLI tool
 - `MILVUS_URI`: Milvus connection URI. **Important**: Do not use quotes around the URI value in your `.env` file (e.g., `MILVUS_URI=http://localhost:19530` instead of `MILVUS_URI="http://localhost:19530"`).
+- `CUSTOM_EMBEDDING_HEADERS`: Custom headers for your embedding provider when using `embedding: custom_local`.
+  **Important**: Due to shell parsing, the value **must be enclosed in single quotes** in your `.env` file to handle special characters correctly.
+  - **Recommended format (JSON string):**
+    ```
+    CUSTOM_EMBEDDING_HEADERS='{"API_SECRET_KEY": "your-secret-key", "Another-Header": "value"}'
+    ```
+  - **Alternative format (key-value pairs):**
+    ```
+    CUSTOM_EMBEDDING_HEADERS='API_SECRET_KEY=your-secret-key,Another-Header=value'
+    ```
 - Database-specific environment variables for Weaviate and Milvus connections
 
 For detailed environment variable usage in CLI and MCP server, see their respective README files.

src/db/vector_db_milvus.py

@@ -113,6 +113,37 @@ def _create_client(self) -> None:
             if original_milvus_uri:
                 os.environ["MILVUS_URI"] = original_milvus_uri
 
+    def _parse_custom_headers(self) -> dict[str, str]:
+        """Parse CUSTOM_EMBEDDING_HEADERS environment variable into a dictionary."""
+        headers_str = os.getenv("CUSTOM_EMBEDDING_HEADERS")
+        if not headers_str:
+            return {}
+
+        # Strip leading/trailing quotes that might come from .env files or shell exports
+        if (headers_str.startswith('"') and headers_str.endswith('"')) or (
+            headers_str.startswith("'") and headers_str.endswith("'")
+        ):
+            headers_str = headers_str[1:-1]
+
+        try:
+            # Try parsing as JSON first
+            headers = json.loads(headers_str)
+            if isinstance(headers, dict):
+                return headers
+            # If JSON parsing results in a non-dict (e.g. a string),
+            # fall through to key-value parsing.
+        except json.JSONDecodeError:
+            # Not a valid JSON object, so fall back to key=value parsing
+            pass
+
+        headers = {}
+        for item in headers_str.split(","):
+            # Split only on the first '=' to allow for '=' in the value
+            key_value = item.split("=", 1)
+            if len(key_value) == 2:
+                headers[key_value[0].strip()] = key_value[1].strip()
+        return headers
+
     def _generate_embedding(self, text: str, embedding_model: str) -> list[float]:
         """
         Generate embeddings for text using the specified model.
@@ -129,7 +160,6 @@ def _generate_embedding(self, text: str, embedding_model: str) -> list[float]:
 
             client_kwargs = {}
             model_to_use = embedding_model
-
             if embedding_model == "custom_local":
                 custom_endpoint_url = os.getenv("CUSTOM_EMBEDDING_URL")
                 if not custom_endpoint_url:
@@ -144,6 +174,11 @@ def _generate_embedding(self, text: str, embedding_model: str) -> list[float]:
                     raise ValueError(
                         "CUSTOM_EMBEDDING_MODEL must be set for 'custom_local' embedding."
                     )
+
+                # Add custom headers if available
+                custom_headers = self._parse_custom_headers()
+                if custom_headers:
+                    client_kwargs["default_headers"] = custom_headers
             else:
                 # Get OpenAI API key from environment
                 api_key = os.getenv("OPENAI_API_KEY")
@@ -157,6 +192,7 @@ def _generate_embedding(self, text: str, embedding_model: str) -> list[float]:
                     model_to_use = "text-embedding-ada-002"
 
             client = openai.OpenAI(**client_kwargs)
+
             response = client.embeddings.create(model=model_to_use, input=text)
 
             return response.data[0].embedding

src/maestro_mcp/README.md

@@ -356,6 +356,16 @@ The server respects the following environment variables:
 - `CUSTOM_EMBEDDING_API_KEY`: The API key for the custom embedding endpoint (optional, but recommended for authentication).
 - `CUSTOM_EMBEDDING_MODEL`: The model name for the custom embedding endpoint (required for `custom_local` embedding for Milvus).
 - `CUSTOM_EMBEDDING_VECTORSIZE`: The vector dimension for the `custom_local` embedding model (required when using `custom_local`).
+- `CUSTOM_EMBEDDING_HEADERS`: Custom headers for your embedding provider when using `embedding: custom_local`.
+  **Important**: Due to shell parsing, the value **must be enclosed in single quotes** in your `.env` file to handle special characters correctly.
+  - **Recommended format (JSON string):**
+    ```
+    CUSTOM_EMBEDDING_HEADERS='{"API_SECRET_KEY": "your-secret-key", "Another-Header": "value"}'
+    ```
+  - **Alternative format (key-value pairs):**
+    ```
+    CUSTOM_EMBEDDING_HEADERS='API_SECRET_KEY=your-secret-key,Another-Header=value'
+    ```
 - Database-specific environment variables for Weaviate and Milvus connections
 
 ## Error Handling

start.sh

@@ -201,7 +201,9 @@ start_http_server() {
     # Load .env file if it exists
     if [ -f "$SCRIPT_DIR/.env" ]; then
         print_status "Loading environment variables from .env file..."
-        export $(grep -v '^#' "$SCRIPT_DIR/.env" | xargs)
+        set -o allexport
+        source "$SCRIPT_DIR/.env"
+        set +o allexport
     fi
 
     if [ -n "$CUSTOM_EMBEDDING_URL" ]; then
@@ -254,7 +256,9 @@ start_stdio_server() {
     # Load .env file if it exists
     if [ -f "$SCRIPT_DIR/.env" ]; then
         print_status "Loading environment variables from .env file..."
-        export $(grep -v '^#' "$SCRIPT_DIR/.env" | xargs)
+        set -o allexport
+        source "$SCRIPT_DIR/.env"
+        set +o allexport
     fi
 
     # Test module import

tests/test_vector_db_milvus.py

@@ -701,6 +701,57 @@ def test_delete_documents_raises_milvus_exception(
         with pytest.raises(MilvusException, match="Delete failed"):
             db.delete_documents(["1"])
 
+    def test_parse_custom_headers(self) -> None:
+        """Test the _parse_custom_headers method with various formats."""
+        db = MilvusVectorDatabase()
+
+        # Test case 1: No environment variable set
+        with patch.dict(os.environ, {}, clear=True):
+            assert db._parse_custom_headers() == {}
+
+        # Test case 2: Simple key-value string
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": "KEY=VALUE"}):
+            assert db._parse_custom_headers() == {"KEY": "VALUE"}
+
+        # Test case 3: Multiple key-value pairs
+        with patch.dict(
+            os.environ, {"CUSTOM_EMBEDDING_HEADERS": "KEY1=VALUE1,KEY2=VALUE2"}
+        ):
+            assert db._parse_custom_headers() == {"KEY1": "VALUE1", "KEY2": "VALUE2"}
+
+        # Test case 4: Key-value with equals sign in value
+        with patch.dict(
+            os.environ, {"CUSTOM_EMBEDDING_HEADERS": "KEY=VALUE_WITH_EQUALS="}
+        ):
+            assert db._parse_custom_headers() == {"KEY": "VALUE_WITH_EQUALS="}
+
+        # Test case 5: Simple JSON string
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": '{"KEY": "VALUE"}'}):
+            assert db._parse_custom_headers() == {"KEY": "VALUE"}
+
+        # Test case 6: JSON string with multiple values
+        with patch.dict(
+            os.environ,
+            {"CUSTOM_EMBEDDING_HEADERS": '{"KEY1": "VALUE1", "KEY2": "VALUE2"}'},
+        ):
+            assert db._parse_custom_headers() == {"KEY1": "VALUE1", "KEY2": "VALUE2"}
+
+        # Test case 7: Key-value string with surrounding quotes to be stripped
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": '"KEY=VALUE"'}):
+            assert db._parse_custom_headers() == {"KEY": "VALUE"}
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": "'KEY=VALUE'"}):
+            assert db._parse_custom_headers() == {"KEY": "VALUE"}
+
+        # Test case 8: Key-value with spaces
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": "  KEY = VALUE  "}):
+            assert db._parse_custom_headers() == {"KEY": "VALUE"}
+
+        # Test case 9: JSON that is not a dictionary (should be ignored)
+        with patch.dict(os.environ, {"CUSTOM_EMBEDDING_HEADERS": '"just a string"'}):
+            # This is valid JSON (a string), but not a dict, so our parser should fall back
+            # to key-value parsing, which will fail and return an empty dict.
+            assert db._parse_custom_headers() == {}
+
 
 if __name__ == "__main__":
     pytest.main()