Upload retriever.py
This commit is contained in:
parent
e29196e458
commit
bd618fd353
41
app/retriever.py
Normal file
41
app/retriever.py
Normal file
@ -0,0 +1,41 @@
|
||||
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
|
||||
Loading…
Reference in New Issue
Block a user