Adding custom connection to weaviate connection client to enable client connection to local container (#3360)
Co-authored-by: Vikram Iyer <vikramiyer@Vikrams-MacBook-Pro.local>
This commit is contained in:
co-authored by
Vikram Iyer
parent
3b2d0ad0eb
commit
8c8368781d
@@ -3,6 +3,7 @@ import uuid
|
|||||||
from typing import Dict, List, Mapping, Optional
|
from typing import Dict, List, Mapping, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import weaviate
|
import weaviate
|
||||||
@@ -12,7 +13,7 @@ except ImportError:
|
|||||||
)
|
)
|
||||||
|
|
||||||
import weaviate.classes.config as wvcc
|
import weaviate.classes.config as wvcc
|
||||||
from weaviate.classes.init import Auth
|
from weaviate.classes.init import Auth, AdditionalConfig, Timeout
|
||||||
from weaviate.classes.query import Filter, MetadataQuery
|
from weaviate.classes.query import Filter, MetadataQuery
|
||||||
from weaviate.util import get_valid_uuid
|
from weaviate.util import get_valid_uuid
|
||||||
|
|
||||||
@@ -47,14 +48,36 @@ class Weaviate(VectorStoreBase):
|
|||||||
auth_config (dict, optional): Authentication configuration for Weaviate. Defaults to None.
|
auth_config (dict, optional): Authentication configuration for Weaviate. Defaults to None.
|
||||||
additional_headers (dict, optional): Additional headers for requests. Defaults to None.
|
additional_headers (dict, optional): Additional headers for requests. Defaults to None.
|
||||||
"""
|
"""
|
||||||
if "localhost" in cluster_url:
|
if "localhost" in cluster_url:
|
||||||
self.client = weaviate.connect_to_local(headers=additional_headers)
|
self.client = weaviate.connect_to_local(headers=additional_headers)
|
||||||
else:
|
elif auth_client_secret:
|
||||||
self.client = weaviate.connect_to_wcs(
|
self.client = weaviate.connect_to_wcs(
|
||||||
cluster_url=cluster_url,
|
cluster_url=cluster_url,
|
||||||
auth_credentials=Auth.api_key(auth_client_secret),
|
auth_credentials=Auth.api_key(auth_client_secret),
|
||||||
headers=additional_headers,
|
headers=additional_headers,
|
||||||
)
|
)
|
||||||
|
else:
|
||||||
|
parsed = urlparse(cluster_url) # e.g., http://mem0_store:8080
|
||||||
|
http_host = parsed.hostname or "localhost"
|
||||||
|
http_port = parsed.port or (443 if parsed.scheme == "https" else 8080)
|
||||||
|
http_secure = parsed.scheme == "https"
|
||||||
|
|
||||||
|
# Weaviate gRPC defaults (inside Docker network)
|
||||||
|
grpc_host = http_host
|
||||||
|
grpc_port = 50051
|
||||||
|
grpc_secure = False
|
||||||
|
|
||||||
|
self.client = weaviate.connect_to_custom(
|
||||||
|
http_host,
|
||||||
|
http_port,
|
||||||
|
http_secure,
|
||||||
|
grpc_host,
|
||||||
|
grpc_port,
|
||||||
|
grpc_secure,
|
||||||
|
headers=additional_headers,
|
||||||
|
skip_init_checks=True,
|
||||||
|
additional_config=AdditionalConfig(timeout=Timeout(init=2.0))
|
||||||
|
)
|
||||||
|
|
||||||
self.collection_name = collection_name
|
self.collection_name = collection_name
|
||||||
self.embedding_model_dims = embedding_model_dims
|
self.embedding_model_dims = embedding_model_dims
|
||||||
|
|||||||
@@ -31,7 +31,6 @@ from fastapi import FastAPI, Request
|
|||||||
from fastapi.routing import APIRouter
|
from fastapi.routing import APIRouter
|
||||||
from mcp.server.fastmcp import FastMCP
|
from mcp.server.fastmcp import FastMCP
|
||||||
from mcp.server.sse import SseServerTransport
|
from mcp.server.sse import SseServerTransport
|
||||||
from qdrant_client import models as qdrant_models
|
|
||||||
|
|
||||||
# Load environment variables
|
# Load environment variables
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
@@ -165,74 +164,54 @@ async def search_memory(query: str) -> str:
|
|||||||
# Get accessible memory IDs based on ACL
|
# Get accessible memory IDs based on ACL
|
||||||
user_memories = db.query(Memory).filter(Memory.user_id == user.id).all()
|
user_memories = db.query(Memory).filter(Memory.user_id == user.id).all()
|
||||||
accessible_memory_ids = [memory.id for memory in user_memories if check_memory_access_permissions(db, memory, app.id)]
|
accessible_memory_ids = [memory.id for memory in user_memories if check_memory_access_permissions(db, memory, app.id)]
|
||||||
|
|
||||||
conditions = [qdrant_models.FieldCondition(key="user_id", match=qdrant_models.MatchValue(value=uid))]
|
|
||||||
|
|
||||||
if accessible_memory_ids:
|
|
||||||
# Convert UUIDs to strings for Qdrant
|
|
||||||
accessible_memory_ids_str = [str(memory_id) for memory_id in accessible_memory_ids]
|
|
||||||
conditions.append(qdrant_models.HasIdCondition(has_id=accessible_memory_ids_str))
|
|
||||||
|
|
||||||
filters = qdrant_models.Filter(must=conditions)
|
filters = {
|
||||||
|
"user_id": uid
|
||||||
|
}
|
||||||
|
|
||||||
embeddings = memory_client.embedding_model.embed(query, "search")
|
embeddings = memory_client.embedding_model.embed(query, "search")
|
||||||
|
|
||||||
hits = memory_client.vector_store.client.query_points(
|
hits = memory_client.vector_store.search(
|
||||||
collection_name=memory_client.vector_store.collection_name,
|
query=query,
|
||||||
query=embeddings,
|
vectors=embeddings,
|
||||||
query_filter=filters,
|
limit=10,
|
||||||
limit=10,
|
filters=filters,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Process search results
|
allowed = set(str(mid) for mid in accessible_memory_ids) if accessible_memory_ids else None
|
||||||
memories = hits.points
|
|
||||||
memories = [
|
|
||||||
{
|
|
||||||
"id": memory.id,
|
|
||||||
"memory": memory.payload["data"],
|
|
||||||
"hash": memory.payload.get("hash"),
|
|
||||||
"created_at": memory.payload.get("created_at"),
|
|
||||||
"updated_at": memory.payload.get("updated_at"),
|
|
||||||
"score": memory.score,
|
|
||||||
}
|
|
||||||
for memory in memories
|
|
||||||
]
|
|
||||||
|
|
||||||
# Log memory access for each memory found
|
results = []
|
||||||
if isinstance(memories, dict) and 'results' in memories:
|
for h in hits:
|
||||||
print(f"Memories: {memories}")
|
# All vector db search functions return OutputData class
|
||||||
for memory_data in memories['results']:
|
id, score, payload = h.id, h.score, h.payload
|
||||||
if 'id' in memory_data:
|
if allowed and h.id is None or h.id not in allowed:
|
||||||
memory_id = uuid.UUID(memory_data['id'])
|
continue
|
||||||
# Create access log entry
|
|
||||||
access_log = MemoryAccessLog(
|
results.append({
|
||||||
memory_id=memory_id,
|
"id": id,
|
||||||
app_id=app.id,
|
"memory": payload.get("data"),
|
||||||
access_type="search",
|
"hash": payload.get("hash"),
|
||||||
metadata_={
|
"created_at": payload.get("created_at"),
|
||||||
"query": query,
|
"updated_at": payload.get("updated_at"),
|
||||||
"score": memory_data.get('score'),
|
"score": score,
|
||||||
"hash": memory_data.get('hash')
|
})
|
||||||
}
|
|
||||||
)
|
for r in results:
|
||||||
db.add(access_log)
|
if r.get("id"):
|
||||||
db.commit()
|
|
||||||
else:
|
|
||||||
for memory in memories:
|
|
||||||
memory_id = uuid.UUID(memory['id'])
|
|
||||||
# Create access log entry
|
|
||||||
access_log = MemoryAccessLog(
|
access_log = MemoryAccessLog(
|
||||||
memory_id=memory_id,
|
memory_id=uuid.UUID(r["id"]),
|
||||||
app_id=app.id,
|
app_id=app.id,
|
||||||
access_type="search",
|
access_type="search",
|
||||||
metadata_={
|
metadata_={
|
||||||
"query": query,
|
"query": query,
|
||||||
"score": memory.get('score'),
|
"score": r.get("score"),
|
||||||
"hash": memory.get('hash')
|
"hash": r.get("hash"),
|
||||||
}
|
},
|
||||||
)
|
)
|
||||||
db.add(access_log)
|
db.add(access_log)
|
||||||
db.commit()
|
db.commit()
|
||||||
return json.dumps(memories, indent=2)
|
|
||||||
|
return json.dumps({"results": results}, indent=2)
|
||||||
finally:
|
finally:
|
||||||
db.close()
|
db.close()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
Reference in New Issue
Block a user