rag-chain-agent/app/retriever.py
2025-04-22 23:56:13 +00:00

42 lines
1.5 KiB
Python

from langchain.vectorstores import Weaviate
from langchain_weaviate.vectorstores import WeaviateVectorStore
import weaviate
from weaviate.classes.init import Auth
from weaviate.config import AdditionalConfig, Timeout
from embedding import WeaviateCustomEmbeddings
import os
http_host = os.getenv("WEAVIATE_HTTP_HOST", "183.111.96.67")
grpc_host = os.getenv("WEAVIATE_GRPC_HOST", "183.111.96.67")
weaviate_api_key = os.getenv("WEAVIATE_API_KEY", "01jse60hwsf37za5kmnkzzcwa9")
def get_retriever():
client = weaviate.connect_to_custom(
http_host=http_host, # Hostname for the HTTP API connection
http_port=32656, # Default is 80, WCD uses 443
http_secure=False, # Whether to use https (secure) for the HTTP API connection
grpc_host=grpc_host, # Hostname for the gRPC API connection
grpc_port=30498, # Default is 50051, WCD uses 443
grpc_secure=False, # Whether to use a secure channel for the gRPC API connection
auth_credentials=Auth.api_key(weaviate_api_key), # API key for authentication
)
embedding = WeaviateCustomEmbeddings()
vectorstore = WeaviateVectorStore(
client=client,
index_name="Loca",
text_key="page_content",
embedding=embedding,
attributes=["source", "page_content"]
)
retriever = vectorstore.as_retriever(
search_kwargs={
"k": 5
}
)
return retriever