traia-iatp 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of traia-iatp might be problematic. Click here for more details.

Files changed (72) hide show
  1. traia_iatp/README.md +368 -0
  2. traia_iatp/__init__.py +30 -0
  3. traia_iatp/cli/__init__.py +5 -0
  4. traia_iatp/cli/main.py +483 -0
  5. traia_iatp/client/__init__.py +10 -0
  6. traia_iatp/client/a2a_client.py +274 -0
  7. traia_iatp/client/crewai_a2a_tools.py +335 -0
  8. traia_iatp/client/grpc_a2a_tools.py +349 -0
  9. traia_iatp/client/root_path_a2a_client.py +1 -0
  10. traia_iatp/core/__init__.py +43 -0
  11. traia_iatp/core/models.py +161 -0
  12. traia_iatp/mcp/__init__.py +15 -0
  13. traia_iatp/mcp/client.py +201 -0
  14. traia_iatp/mcp/mcp_agent_template.py +422 -0
  15. traia_iatp/mcp/templates/Dockerfile.j2 +56 -0
  16. traia_iatp/mcp/templates/README.md.j2 +212 -0
  17. traia_iatp/mcp/templates/cursor-rules.md.j2 +326 -0
  18. traia_iatp/mcp/templates/deployment_params.json.j2 +20 -0
  19. traia_iatp/mcp/templates/docker-compose.yml.j2 +23 -0
  20. traia_iatp/mcp/templates/dockerignore.j2 +47 -0
  21. traia_iatp/mcp/templates/gitignore.j2 +77 -0
  22. traia_iatp/mcp/templates/mcp_health_check.py.j2 +150 -0
  23. traia_iatp/mcp/templates/pyproject.toml.j2 +26 -0
  24. traia_iatp/mcp/templates/run_local_docker.sh.j2 +94 -0
  25. traia_iatp/mcp/templates/server.py.j2 +240 -0
  26. traia_iatp/mcp/traia_mcp_adapter.py +381 -0
  27. traia_iatp/preview_diagrams.html +181 -0
  28. traia_iatp/registry/__init__.py +26 -0
  29. traia_iatp/registry/atlas_search_indexes.json +280 -0
  30. traia_iatp/registry/embeddings.py +298 -0
  31. traia_iatp/registry/iatp_search_api.py +839 -0
  32. traia_iatp/registry/mongodb_registry.py +771 -0
  33. traia_iatp/registry/readmes/ATLAS_SEARCH_INDEXES.md +252 -0
  34. traia_iatp/registry/readmes/ATLAS_SEARCH_SETUP.md +134 -0
  35. traia_iatp/registry/readmes/AUTHENTICATION_UPDATE.md +124 -0
  36. traia_iatp/registry/readmes/EMBEDDINGS_SETUP.md +172 -0
  37. traia_iatp/registry/readmes/IATP_SEARCH_API_GUIDE.md +257 -0
  38. traia_iatp/registry/readmes/MONGODB_X509_AUTH.md +208 -0
  39. traia_iatp/registry/readmes/README.md +251 -0
  40. traia_iatp/registry/readmes/REFACTORING_SUMMARY.md +191 -0
  41. traia_iatp/server/__init__.py +15 -0
  42. traia_iatp/server/a2a_server.py +215 -0
  43. traia_iatp/server/example_template_usage.py +72 -0
  44. traia_iatp/server/iatp_server_agent_generator.py +237 -0
  45. traia_iatp/server/iatp_server_template_generator.py +235 -0
  46. traia_iatp/server/templates/Dockerfile.j2 +49 -0
  47. traia_iatp/server/templates/README.md +137 -0
  48. traia_iatp/server/templates/README.md.j2 +425 -0
  49. traia_iatp/server/templates/__init__.py +1 -0
  50. traia_iatp/server/templates/__main__.py.j2 +450 -0
  51. traia_iatp/server/templates/agent.py.j2 +80 -0
  52. traia_iatp/server/templates/agent_config.json.j2 +22 -0
  53. traia_iatp/server/templates/agent_executor.py.j2 +264 -0
  54. traia_iatp/server/templates/docker-compose.yml.j2 +23 -0
  55. traia_iatp/server/templates/env.example.j2 +67 -0
  56. traia_iatp/server/templates/gitignore.j2 +78 -0
  57. traia_iatp/server/templates/grpc_server.py.j2 +218 -0
  58. traia_iatp/server/templates/pyproject.toml.j2 +76 -0
  59. traia_iatp/server/templates/run_local_docker.sh.j2 +103 -0
  60. traia_iatp/server/templates/server.py.j2 +190 -0
  61. traia_iatp/special_agencies/__init__.py +4 -0
  62. traia_iatp/special_agencies/registry_search_agency.py +392 -0
  63. traia_iatp/utils/__init__.py +10 -0
  64. traia_iatp/utils/docker_utils.py +251 -0
  65. traia_iatp/utils/general.py +64 -0
  66. traia_iatp/utils/iatp_utils.py +126 -0
  67. traia_iatp-0.1.1.dist-info/METADATA +414 -0
  68. traia_iatp-0.1.1.dist-info/RECORD +72 -0
  69. traia_iatp-0.1.1.dist-info/WHEEL +5 -0
  70. traia_iatp-0.1.1.dist-info/entry_points.txt +2 -0
  71. traia_iatp-0.1.1.dist-info/licenses/LICENSE +21 -0
  72. traia_iatp-0.1.1.dist-info/top_level.txt +1 -0
@@ -0,0 +1,252 @@
1
+ # MongoDB Atlas Search Indexes for IATP Registry
2
+
3
+ This document describes the MongoDB Atlas Search and Vector Search indexes required for the IATP utility agent registry.
4
+
5
+ ## Prerequisites
6
+
7
+ 1. MongoDB Atlas cluster (M10 or higher for Vector Search)
8
+ 2. Access to Atlas UI or Atlas Admin API
9
+ 3. OpenAI API key for generating embeddings (set as `OPENAI_API_KEY` environment variable)
10
+
11
+ ## Index Configuration
12
+
13
+ ### 1. Utility Agent Atlas Search Index
14
+
15
+ Create this index on the `utility_agents` collection (or `utility_agents_test`/`utility_agents_prod` based on environment).
16
+
17
+ **Index Name**: `utility_agent_atlas_search` (or `utility_agent_atlas_search_test`/`utility_agent_atlas_search_prod`)
18
+
19
+ **Index Definition**:
20
+ ```json
21
+ {
22
+ "mappings": {
23
+ "dynamic": false,
24
+ "fields": {
25
+ "name": {
26
+ "type": "string",
27
+ "analyzer": "lucene.standard"
28
+ },
29
+ "description": {
30
+ "type": "string",
31
+ "analyzer": "lucene.standard"
32
+ },
33
+ "tags": {
34
+ "type": "string",
35
+ "analyzer": "lucene.standard"
36
+ },
37
+ "capabilities": {
38
+ "type": "string",
39
+ "analyzer": "lucene.standard"
40
+ },
41
+ "search_text": {
42
+ "type": "string",
43
+ "analyzer": "lucene.standard"
44
+ },
45
+ "skills": {
46
+ "type": "document",
47
+ "fields": {
48
+ "name": {
49
+ "type": "string",
50
+ "analyzer": "lucene.standard"
51
+ },
52
+ "description": {
53
+ "type": "string",
54
+ "analyzer": "lucene.standard"
55
+ },
56
+ "examples": {
57
+ "type": "string",
58
+ "analyzer": "lucene.standard"
59
+ }
60
+ }
61
+ },
62
+ "agent_card": {
63
+ "type": "document",
64
+ "fields": {
65
+ "name": {
66
+ "type": "string",
67
+ "analyzer": "lucene.standard"
68
+ },
69
+ "description": {
70
+ "type": "string",
71
+ "analyzer": "lucene.standard"
72
+ }
73
+ }
74
+ }
75
+ }
76
+ }
77
+ }
78
+ ```
79
+
80
+ ### 2. Utility Agent Vector Search Index
81
+
82
+ Create this index for semantic search using embeddings.
83
+
84
+ **Index Name**: `utility_agent_vector_search` (or `utility_agent_vector_search_test`/`utility_agent_vector_search_prod`)
85
+
86
+ **Index Definition**:
87
+ ```json
88
+ {
89
+ "type": "vectorSearch",
90
+ "fields": [
91
+ {
92
+ "type": "vector",
93
+ "path": "embeddings.search_text",
94
+ "numDimensions": 1536,
95
+ "similarity": "cosine"
96
+ },
97
+ {
98
+ "type": "vector",
99
+ "path": "embeddings.description",
100
+ "numDimensions": 1536,
101
+ "similarity": "cosine"
102
+ },
103
+ {
104
+ "type": "vector",
105
+ "path": "embeddings.capabilities",
106
+ "numDimensions": 1536,
107
+ "similarity": "cosine"
108
+ },
109
+ {
110
+ "type": "vector",
111
+ "path": "embeddings.skills",
112
+ "numDimensions": 1536,
113
+ "similarity": "cosine"
114
+ },
115
+ {
116
+ "type": "filter",
117
+ "path": "is_active"
118
+ },
119
+ {
120
+ "type": "filter",
121
+ "path": "tags"
122
+ }
123
+ ]
124
+ }
125
+ ```
126
+
127
+ ### 3. MCP Server Atlas Search Index
128
+
129
+ Create this index on the `mcp_servers` collection.
130
+
131
+ **Index Name**: `mcp_server_atlas_search` (or `mcp_server_atlas_search_test`/`mcp_server_atlas_search_prod`)
132
+
133
+ **Index Definition**:
134
+ ```json
135
+ {
136
+ "mappings": {
137
+ "dynamic": false,
138
+ "fields": {
139
+ "name": {
140
+ "type": "string",
141
+ "analyzer": "lucene.standard"
142
+ },
143
+ "description": {
144
+ "type": "string",
145
+ "analyzer": "lucene.standard"
146
+ },
147
+ "capabilities": {
148
+ "type": "string",
149
+ "analyzer": "lucene.standard"
150
+ }
151
+ }
152
+ }
153
+ }
154
+ ```
155
+
156
+ ### 4. MCP Server Vector Search Index
157
+
158
+ **Index Name**: `mcp_server_vector_search` (or `mcp_server_vector_search_test`/`mcp_server_vector_search_prod`)
159
+
160
+ **Index Definition**:
161
+ ```json
162
+ {
163
+ "type": "vectorSearch",
164
+ "fields": [
165
+ {
166
+ "type": "vector",
167
+ "path": "description_embedding",
168
+ "numDimensions": 1536,
169
+ "similarity": "cosine"
170
+ },
171
+ {
172
+ "type": "vector",
173
+ "path": "capabilities_embedding",
174
+ "numDimensions": 1536,
175
+ "similarity": "cosine"
176
+ },
177
+ {
178
+ "type": "filter",
179
+ "path": "is_active"
180
+ }
181
+ ]
182
+ }
183
+ ```
184
+
185
+ ## Creating Indexes
186
+
187
+ ### Via Atlas UI
188
+
189
+ 1. Navigate to your cluster in MongoDB Atlas
190
+ 2. Go to "Search" tab
191
+ 3. Click "Create Search Index"
192
+ 4. Choose "JSON Editor"
193
+ 5. Select the appropriate database and collection
194
+ 6. Paste the index definition
195
+ 7. Name the index according to the pattern above
196
+ 8. Click "Create Search Index"
197
+
198
+ ### Via Atlas Admin API
199
+
200
+ ```bash
201
+ # Set your API credentials
202
+ export ATLAS_PUBLIC_KEY="your-public-key"
203
+ export ATLAS_PRIVATE_KEY="your-private-key"
204
+ export ATLAS_PROJECT_ID="your-project-id"
205
+ export ATLAS_CLUSTER_NAME="your-cluster-name"
206
+
207
+ # Create utility agent Atlas Search index
208
+ curl --user "${ATLAS_PUBLIC_KEY}:${ATLAS_PRIVATE_KEY}" --digest \
209
+ --header "Content-Type: application/json" \
210
+ --request POST "https://cloud.mongodb.com/api/atlas/v2/groups/${ATLAS_PROJECT_ID}/clusters/${ATLAS_CLUSTER_NAME}/search/indexes" \
211
+ --data '{
212
+ "collectionName": "utility_agents",
213
+ "database": "iatp",
214
+ "name": "utility_agent_atlas_search",
215
+ "mappings": { ... }
216
+ }'
217
+ ```
218
+
219
+ ## Index Naming Convention
220
+
221
+ The indexes follow environment-specific naming:
222
+ - Test environment: `*_test` suffix
223
+ - Production environment: `*_prod` suffix
224
+ - Default/development: No suffix
225
+
226
+ ## Required Embedding Dimensions
227
+
228
+ All vector fields use OpenAI's `text-embedding-3-small` model which produces 1536-dimensional vectors.
229
+
230
+ ## Monitoring Index Creation
231
+
232
+ Index creation can take several minutes. Monitor progress:
233
+ 1. In Atlas UI: Check the "Search" tab for index status
234
+ 2. Via API: Query the index endpoint to check status
235
+
236
+ ## Testing Indexes
237
+
238
+ After creation, test the indexes using the registry methods:
239
+
240
+ ```python
241
+ # Test Atlas Search
242
+ results = await registry.atlas_search("trading hyperliquid")
243
+
244
+ # Test Vector Search
245
+ results = await registry.vector_search_text("find me a trading agent for crypto")
246
+ ```
247
+
248
+ ## Maintenance
249
+
250
+ - Indexes are automatically maintained by Atlas
251
+ - No manual optimization required
252
+ - Monitor index performance in Atlas UI under "Search Metrics"
@@ -0,0 +1,134 @@
1
+ # MongoDB Atlas Search Index Setup
2
+
3
+ This guide explains how to set up Atlas Search and Vector Search indexes for the IATP registry collections.
4
+
5
+ ## Overview
6
+
7
+ Each collection requires two types of indexes:
8
+ 1. **Atlas Search Index** - For text-based keyword search
9
+ 2. **Vector Search Index** - For semantic similarity search using embeddings
10
+
11
+ ## Index Naming Convention
12
+
13
+ All index names include the environment suffix:
14
+ - `utility_agency_atlas_search_{env}`
15
+ - `utility_agency_vector_search_{env}`
16
+ - `mcp_server_atlas_search_{env}`
17
+ - `mcp_server_vector_search_{env}`
18
+
19
+ Where `{env}` is one of: `test`, `staging`, or `prod`
20
+
21
+ ## Collections
22
+
23
+ Based on your environment (ENV variable), the collections are:
24
+ - `iatp-utility-agency-registry-{env}` (where env = test/staging/prod)
25
+ - `iatp-mcp-server-registry-{env}`
26
+
27
+ ## Creating Indexes in MongoDB Atlas
28
+
29
+ ### Method 1: Atlas UI
30
+
31
+ 1. Go to your MongoDB Atlas cluster
32
+ 2. Click on "Search" in the left sidebar
33
+ 3. Click "Create Search Index"
34
+ 4. Choose your database (`iatp`)
35
+ 5. For each collection, create both indexes:
36
+
37
+ #### Utility Agency Registry - Atlas Search Index
38
+ - Collection: `iatp-utility-agency-registry-{env}`
39
+ - Index Name: `utility_agency_atlas_search_{env}`
40
+ - Use the JSON definition from `atlas_search_indexes.json` under `utility_agency_indexes.atlas_search`
41
+
42
+ #### Utility Agency Registry - Vector Search Index
43
+ - Collection: `iatp-utility-agency-registry-{env}`
44
+ - Index Name: `utility_agency_vector_search_{env}`
45
+ - Use the JSON definition from `atlas_search_indexes.json` under `utility_agency_indexes.vector_search`
46
+
47
+ #### MCP Server Registry - Atlas Search Index
48
+ - Collection: `iatp-mcp-server-registry-{env}`
49
+ - Index Name: `mcp_server_atlas_search_{env}`
50
+ - Use the JSON definition from `atlas_search_indexes.json` under `mcp_server_indexes.atlas_search`
51
+
52
+ #### MCP Server Registry - Vector Search Index
53
+ - Collection: `iatp-mcp-server-registry-{env}`
54
+ - Index Name: `mcp_server_vector_search_{env}`
55
+ - Use the JSON definition from `atlas_search_indexes.json` under `mcp_server_indexes.vector_search`
56
+
57
+ ### Method 2: Atlas Admin API
58
+
59
+ ```bash
60
+ # Set your Atlas API credentials
61
+ export ATLAS_PUBLIC_KEY="your_public_key"
62
+ export ATLAS_PRIVATE_KEY="your_private_key"
63
+ export ATLAS_PROJECT_ID="your_project_id"
64
+ export ATLAS_CLUSTER_NAME="your_cluster_name"
65
+ export ENV="test" # or staging, prod
66
+
67
+ # Create indexes using curl (example for one index)
68
+ # Note: You'll need to modify the JSON to replace <env> with your actual environment
69
+ curl --user "${ATLAS_PUBLIC_KEY}:${ATLAS_PRIVATE_KEY}" --digest \
70
+ --header "Content-Type: application/json" \
71
+ --request POST "https://cloud.mongodb.com/api/atlas/v1.0/groups/${ATLAS_PROJECT_ID}/clusters/${ATLAS_CLUSTER_NAME}/fts/indexes" \
72
+ --data @atlas_search_indexes.json
73
+ ```
74
+
75
+ ## Example: Creating Indexes for Test Environment
76
+
77
+ For the `test` environment, you would create:
78
+ - `utility_agency_atlas_search_test`
79
+ - `utility_agency_vector_search_test`
80
+ - `mcp_server_atlas_search_test`
81
+ - `mcp_server_vector_search_test`
82
+
83
+ ## Vector Embeddings
84
+
85
+ The vector search indexes expect the following embedding fields:
86
+ - `description_embedding` - 1536-dimensional vector for description text
87
+ - `tags_embedding` - 1536-dimensional vector for tags (utility agencies)
88
+ - `capabilities_embedding` - 1536-dimensional vector for capabilities (MCP servers)
89
+
90
+ ### Generating Embeddings
91
+
92
+ To use vector search, you need to generate embeddings when adding documents. Example using OpenAI:
93
+
94
+ ```python
95
+ import openai
96
+
97
+ def generate_embedding(text: str) -> List[float]:
98
+ response = openai.Embedding.create(
99
+ model="text-embedding-ada-002",
100
+ input=text
101
+ )
102
+ return response['data'][0]['embedding']
103
+ ```
104
+
105
+ ## Using the Indexes
106
+
107
+ The code automatically uses the correct index name based on your ENV variable:
108
+
109
+ ### Atlas Search (Text Search)
110
+ ```python
111
+ # The atlas_search method automatically uses the correct index name
112
+ results = await registry.atlas_search("weather forecast")
113
+ ```
114
+
115
+ ### Vector Search (Semantic Search)
116
+ ```python
117
+ # The vector_search method automatically uses the correct index name
118
+ query_embedding = generate_embedding("meteorological predictions")
119
+ results = await registry.vector_search(query_embedding)
120
+ ```
121
+
122
+ ## Important Notes
123
+
124
+ 1. **Index Status**: After creating indexes, they may take a few minutes to build
125
+ 2. **Environment Variable**: Make sure your `ENV` environment variable is set correctly
126
+ 3. **Embedding Dimensions**: The vector indexes are configured for 1536 dimensions (OpenAI's ada-002 model)
127
+ 4. **Similarity Metric**: Using cosine similarity for vector comparison
128
+ 5. **Filters**: The vector indexes include `is_active` as a filter field for efficient filtering
129
+
130
+ ## Monitoring
131
+
132
+ You can monitor index status and usage in the Atlas UI under:
133
+ - Search > Index Status
134
+ - Search > Query Analytics
@@ -0,0 +1,124 @@
1
+ # MongoDB Authentication Update Summary
2
+
3
+ ## Overview
4
+
5
+ We have simplified MongoDB authentication by removing API key support and focusing on two primary authentication methods:
6
+
7
+ 1. **X.509 Certificate Authentication** (Recommended)
8
+ 2. **Username/Password Authentication** (Fallback)
9
+
10
+ ## Changes Made
11
+
12
+ ### 1. Code Updates
13
+
14
+ #### mongodb_registry.py
15
+ - Removed all references to `MONGODB_READWRITE_API_KEY` and `MONGODB_READONLY_API_KEY`
16
+ - Updated both `UtilityAgentRegistry` and `MCPServerRegistry` classes
17
+ - Simplified authentication priority to: Certificate → Username/Password → Connection String
18
+
19
+ #### iatp_registry_api.py
20
+ - Removed API key authentication from `get_readonly_connection_string()`
21
+ - Updated error messages to reflect new authentication methods
22
+
23
+ ### 2. Removed Files
24
+ - `mongodb_data_api_client.py` - MongoDB Data API client (API key based)
25
+ - `update_env_for_api_keys.py` - API key setup utility
26
+ - `test_end_to_end_with_api_keys.py` - API key specific test
27
+
28
+ ### 3. Documentation Updates
29
+ - Updated `MONGODB_X509_AUTH.md` to reflect simplified authentication priority
30
+ - Updated `README.md` files to remove API key references
31
+ - Updated `search_api_service.py` documentation
32
+
33
+ ### 4. Test Updates
34
+ - Updated `test_mongodb_x509_auth.py` to remove API key checks
35
+ - Created `test_end_to_end_with_auth.py` as replacement for API key test
36
+
37
+ ## Authentication Methods
38
+
39
+ ### Primary: X.509 Certificate
40
+ ```bash
41
+ export MONGODB_X509_CERT_FILE=/path/to/certificate.pem
42
+ ```
43
+
44
+ **Benefits:**
45
+ - Most secure - no passwords in environment
46
+ - Certificate-based mutual TLS authentication
47
+ - Can be rotated without code changes
48
+ - Supported by MongoDB Atlas
49
+
50
+ ### Fallback: Username/Password
51
+ ```bash
52
+ export MONGODB_USER=username
53
+ export MONGODB_PASSWORD=password
54
+ ```
55
+
56
+ **When to use:**
57
+ - Development environments
58
+ - Legacy systems
59
+ - When certificates are not available
60
+
61
+ ### Connection String
62
+ ```bash
63
+ export MONGODB_CONNECTION_STRING=mongodb+srv://...
64
+ ```
65
+
66
+ **When to use:**
67
+ - Complex connection configurations
68
+ - Custom authentication mechanisms
69
+
70
+ ## Migration Guide
71
+
72
+ ### From API Keys to X.509 Certificates
73
+
74
+ 1. Create X.509 user in MongoDB Atlas
75
+ 2. Download certificate file (.pem)
76
+ 3. Update environment:
77
+ ```bash
78
+ # Remove
79
+ unset MONGODB_READWRITE_API_KEY
80
+ unset MONGODB_READONLY_API_KEY
81
+
82
+ # Add
83
+ export MONGODB_X509_CERT_FILE=/path/to/cert.pem
84
+ ```
85
+ 4. Test with: `uv run python tests/test_mongodb_x509_auth.py`
86
+
87
+ ### From API Keys to Username/Password
88
+
89
+ 1. Update environment:
90
+ ```bash
91
+ # Remove
92
+ unset MONGODB_READWRITE_API_KEY
93
+ unset MONGODB_READONLY_API_KEY
94
+
95
+ # Add
96
+ export MONGODB_USER=your-username
97
+ export MONGODB_PASSWORD=your-password
98
+ ```
99
+
100
+ ## Testing
101
+
102
+ Run authentication tests:
103
+ ```bash
104
+ # Test X.509 authentication
105
+ uv run python tests/test_mongodb_registry/test_mongodb_x509_auth.py
106
+
107
+ # Test end-to-end with current auth method
108
+ uv run python tests/test_mongodb_registry/test_end_to_end_with_auth.py
109
+ ```
110
+
111
+ ## Security Recommendations
112
+
113
+ 1. **Use X.509 certificates in production** - Most secure option
114
+ 2. **Protect certificate files** - Set file permissions to 600
115
+ 3. **Rotate credentials regularly** - Both certificates and passwords
116
+ 4. **Use environment-specific credentials** - Different for dev/staging/prod
117
+ 5. **Never commit credentials** - Use secure secret management
118
+
119
+ ## Backward Compatibility
120
+
121
+ The system maintains backward compatibility:
122
+ - Existing username/password authentication continues to work
123
+ - Connection strings are still supported
124
+ - The authentication priority ensures smooth transition
@@ -0,0 +1,172 @@
1
+ # Embeddings Setup for Vector Search
2
+
3
+ This guide explains how to enable automatic embedding generation for vector search in IATP.
4
+
5
+ ## Important Note
6
+
7
+ MongoDB Atlas does **not** automatically generate embeddings. The IATP application handles embedding generation:
8
+ - When documents are added, IATP generates embeddings using your chosen provider (OpenAI, Cohere, etc.)
9
+ - When searching, IATP converts your text query to embeddings before sending to MongoDB
10
+ - This happens transparently when `ENABLE_EMBEDDINGS=true`
11
+
12
+ ## Overview
13
+
14
+ IATP can automatically generate embeddings when adding documents to the registry, enabling semantic vector search. This is optional - the system works without embeddings, but you won't be able to use vector search.
15
+
16
+ ## Environment Variables
17
+
18
+ To enable embeddings, set the following environment variables:
19
+
20
+ ```bash
21
+ # Enable embedding generation
22
+ export ENABLE_EMBEDDINGS=true
23
+
24
+ # Choose your embedding provider (default: openai)
25
+ # Options: openai, cohere
26
+ export EMBEDDING_PROVIDER=openai
27
+
28
+ # Provider-specific API keys
29
+ export OPENAI_API_KEY=your_openai_api_key # For OpenAI
30
+ export COHERE_API_KEY=your_cohere_api_key # For Cohere
31
+ ```
32
+
33
+ ## How It Works
34
+
35
+ ### 1. Document Storage with Embeddings
36
+
37
+ When `ENABLE_EMBEDDINGS=true`, the system automatically generates embeddings:
38
+
39
+ **For Utility Agencies:**
40
+ - `description_embedding`: Vector embedding of the description text
41
+ - `tags_embedding`: Vector embedding of concatenated tags
42
+
43
+ **For MCP Servers:**
44
+ - `description_embedding`: Vector embedding of the description text
45
+ - `capabilities_embedding`: Vector embedding of concatenated capabilities
46
+
47
+ ### 2. Search Methods
48
+
49
+ **Without Embeddings (Always Available):**
50
+ ```python
51
+ # Basic keyword search (requires MongoDB text index)
52
+ results = await registry.query_agencies(query="weather")
53
+
54
+ # Atlas Search (requires Atlas Search index)
55
+ results = await registry.atlas_search("weather forecast")
56
+ ```
57
+
58
+ **With Embeddings (Requires ENABLE_EMBEDDINGS=true):**
59
+ ```python
60
+ # Vector search with text query (auto-converts to embedding)
61
+ results = await registry.vector_search_text("meteorological predictions")
62
+
63
+ # Vector search with specific field
64
+ results = await registry.vector_search_text("api integration", search_field="tags")
65
+
66
+ # Vector search with pre-computed embedding
67
+ embedding = await embedding_service.generate_embedding("weather")
68
+ results = await registry.vector_search(embedding)
69
+ ```
70
+
71
+ ## Embedding Providers
72
+
73
+ ### OpenAI (Default)
74
+ - Model: `text-embedding-ada-002` (1536 dimensions)
75
+ - Best for: General-purpose semantic search
76
+ - Cost: ~$0.0001 per 1K tokens
77
+
78
+ ### Cohere
79
+ - Model: `embed-english-v3.0` (1024 dimensions)
80
+ - Best for: Multilingual support
81
+ - Note: Update vector index dimensions to 1024 if using Cohere
82
+
83
+ ## Installation
84
+
85
+ IATP includes optional dependencies for embeddings support. Install using uv:
86
+
87
+ ```bash
88
+ # Install IATP with embeddings support (includes all providers)
89
+ uv sync --extra embeddings
90
+
91
+ # Or add specific providers to your project
92
+ uv add openai # For OpenAI embeddings
93
+ uv add cohere # For Cohere embeddings
94
+ ```
95
+
96
+ The embedding providers are defined as optional dependencies in `pyproject.toml`:
97
+ ```toml
98
+ [project.optional-dependencies]
99
+ embeddings = [
100
+ "openai>=1.0.0",
101
+ "cohere>=5.0.0",
102
+ ]
103
+ ```
104
+
105
+ ### Development Note
106
+
107
+ When developing with uv:
108
+ - Use `uv sync --extra embeddings` to install optional dependencies
109
+ - Use `uv add <package>` to add new dependencies to your project
110
+ - Dependencies added with `uv add` will be added to the main dependencies in `pyproject.toml`
111
+
112
+ ## Example Usage
113
+
114
+ ```python
115
+ import os
116
+ os.environ["ENABLE_EMBEDDINGS"] = "true"
117
+ os.environ["OPENAI_API_KEY"] = "your-key"
118
+
119
+ from iatp.registry.mongodb_registry import UtilityAgencyRegistry
120
+
121
+ # Create registry (embeddings will be auto-generated)
122
+ registry = UtilityAgencyRegistry()
123
+
124
+ # Add agency - embeddings generated automatically
125
+ await registry.add_agency(agency, endpoint, tags=["weather", "api"])
126
+
127
+ # Search using semantic similarity
128
+ results = await registry.vector_search_text("climate data analysis")
129
+ ```
130
+
131
+ ## Performance Considerations
132
+
133
+ 1. **Embedding Generation Time**: ~100-500ms per text
134
+ 2. **Batch Processing**: Use `generate_embeddings()` for multiple texts
135
+ 3. **Caching**: Consider caching embeddings for frequently searched queries
136
+ 4. **Cost**: Monitor API usage, especially for large datasets
137
+
138
+ ## Troubleshooting
139
+
140
+ ### "Failed to generate embeddings" Warning
141
+ - Check that `ENABLE_EMBEDDINGS=true`
142
+ - Verify API key is set correctly
143
+ - Check network connectivity
144
+ - Verify provider package is installed
145
+
146
+ ### Vector Search Returns No Results
147
+ - Ensure documents have embeddings (added after enabling embeddings)
148
+ - Verify Vector Search index exists in Atlas
149
+ - Check that embedding dimensions match (1536 for OpenAI)
150
+
151
+ ### Re-indexing Existing Data
152
+ If you have existing data without embeddings:
153
+
154
+ ```python
155
+ # Script to add embeddings to existing documents
156
+ async def reindex_with_embeddings():
157
+ registry = UtilityAgencyRegistry()
158
+
159
+ # Get all agencies
160
+ all_agencies = await registry.query_agencies(limit=1000)
161
+
162
+ for agency in all_agencies:
163
+ # Re-add to generate embeddings
164
+ await registry.add_agency(agency, agency.endpoint, agency.tags)
165
+ ```
166
+
167
+ ## Security Notes
168
+
169
+ 1. **API Keys**: Never commit API keys to version control
170
+ 2. **Rate Limits**: Implement rate limiting for production
171
+ 3. **Error Handling**: System continues without embeddings if generation fails
172
+ 4. **Data Privacy**: Be aware that text is sent to external APIs for embedding