Skip to content

Commit 8a99ea2

Browse files
feat: add framework-agnostic weaviate tool
Co-Authored-By: root@a10k.co <root@a10k.co>
1 parent 5da4248 commit 8a99ea2

2 files changed

Lines changed: 118 additions & 0 deletions

File tree

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import os
2+
import json
3+
import weaviate
4+
from typing import Optional
5+
from weaviate.classes.config import Configure
6+
from weaviate.classes.init import Auth
7+
8+
def search_collection(
9+
collection_name: str,
10+
query: str,
11+
limit: int = 3,
12+
model: str = "nomic-embed-text"
13+
) -> str:
14+
"""Search a Weaviate collection using near-text queries.
15+
16+
Args:
17+
collection_name: Name of the collection to search
18+
query: The search query
19+
limit: Maximum number of results (default: 3)
20+
model: Text embedding model to use (default: nomic-embed-text)
21+
22+
Returns:
23+
str: JSON string containing search results
24+
"""
25+
url = os.environ.get("WEAVIATE_TOOL_URL")
26+
api_key = os.environ.get("WEAVIATE_TOOL_API_KEY")
27+
openai_key = os.environ.get("WEAVIATE_TOOL_OPENAI_API_KEY")
28+
29+
if not url or not api_key or not openai_key:
30+
raise ValueError("Missing required environment variables")
31+
32+
headers = {"X-OpenAI-Api-Key": openai_key}
33+
vectorizer = Configure.Vectorizer.text2vec_openai(model=model)
34+
35+
client = weaviate.connect_to_weaviate_cloud(
36+
cluster_url=url,
37+
auth_credentials=Auth.api_key(api_key),
38+
headers=headers
39+
)
40+
41+
try:
42+
collection = client.collections.get(collection_name)
43+
if not collection:
44+
raise ValueError(f"Collection {collection_name} not found")
45+
46+
response = collection.query.near_text(
47+
query=query,
48+
limit=limit
49+
)
50+
51+
results = []
52+
for obj in response.objects:
53+
results.append(obj.properties)
54+
55+
return json.dumps(results, indent=2)
56+
finally:
57+
client.close()
58+
59+
def create_collection(
60+
collection_name: str,
61+
model: str = "nomic-embed-text"
62+
) -> str:
63+
"""Create a new Weaviate collection.
64+
65+
Args:
66+
collection_name: Name of the collection to create
67+
model: Text embedding model to use (default: nomic-embed-text)
68+
69+
Returns:
70+
str: Success message
71+
"""
72+
url = os.environ.get("WEAVIATE_TOOL_URL")
73+
api_key = os.environ.get("WEAVIATE_TOOL_API_KEY")
74+
openai_key = os.environ.get("WEAVIATE_TOOL_OPENAI_API_KEY")
75+
76+
if not url or not api_key or not openai_key:
77+
raise ValueError("Missing required environment variables")
78+
79+
headers = {"X-OpenAI-Api-Key": openai_key}
80+
vectorizer = Configure.Vectorizer.text2vec_openai(model=model)
81+
82+
client = weaviate.connect_to_weaviate_cloud(
83+
cluster_url=url,
84+
auth_credentials=Auth.api_key(api_key),
85+
headers=headers
86+
)
87+
88+
try:
89+
collection = client.collections.get(collection_name)
90+
if collection:
91+
return f"Collection {collection_name} already exists"
92+
93+
client.collections.create(
94+
name=collection_name,
95+
vectorizer_config=vectorizer
96+
)
97+
return f"Created collection {collection_name}"
98+
finally:
99+
client.close()
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
{
2+
"name": "weaviate",
3+
"url": "https://github.com/weaviate/weaviate-python-client",
4+
"category": "vector-store",
5+
"env": {
6+
"WEAVIATE_TOOL_URL": null,
7+
"WEAVIATE_TOOL_API_KEY": null,
8+
"WEAVIATE_TOOL_OPENAI_API_KEY": null
9+
},
10+
"dependencies": [
11+
"weaviate-client>=3.0.0",
12+
"openai>=1.0.0"
13+
],
14+
"tools": [
15+
"search_collection",
16+
"create_collection"
17+
],
18+
"cta": "🔗 Create your Weaviate cluster here: https://console.weaviate.cloud/"
19+
}

0 commit comments

Comments
 (0)