Skip to content

Commit

Permalink
Switch default vector store to Weaviate
Browse files Browse the repository at this point in the history
  • Loading branch information
@patel-lyzr
patel-lyzr committed Jan 11, 2024
1 parent 5672663 commit 35440e7
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 24 deletions.
39 changes: 28 additions & 11 deletions lyzr/base/vector_store.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
from typing import Optional, Sequence

import os
import uuid
import weaviate
from weaviate.embedded import EmbeddedOptions
from llama_index import Document, ServiceContext, VectorStoreIndex, StorageContext
from llama_index.node_parser import SimpleNodeParser

Expand All @@ -13,30 +17,43 @@ def import_vector_store_class(vector_store_class_name: str):
class LyzrVectorStoreIndex:
@staticmethod
def from_defaults(
vector_store_type: str = "LanceDBVectorStore",
vector_store_type: str = "WeaviateVectorStore",
documents: Optional[Sequence[Document]] = None,
service_context: Optional[ServiceContext] = None,
**kwargs
) -> VectorStoreIndex:
if documents is None and vector_store_type == "SimpleVectorStore":
raise ValueError("documents must be provided for SimpleVectorStore")

vector_store_class = import_vector_store_class(vector_store_type)
VectorStoreClass = import_vector_store_class(vector_store_type)

if vector_store_type == "WeaviateVectorStore":
weaviate_client = weaviate.Client(
embedded_options=weaviate.embedded.EmbeddedOptions(),
additional_headers={"X-OpenAI-Api-Key": os.environ["OPENAI_API_KEY"]},
)
kwargs["weaviate_client"] = (
weaviate_client
if "weaviate_client" not in kwargs
else kwargs["weaviate_client"]
)
kwargs["index_name"] = (
uuid if "index_name" not in kwargs else kwargs["index_name"]
)

vector_store = VectorStoreClass(**kwargs)
else:
vector_store = VectorStoreClass(**kwargs)

if documents is None:
vector_store = vector_store_class(**kwargs)
index = VectorStoreIndex.from_vector_store(
vector_store=vector_store, service_context=service_context
)
else:
if vector_store_type == "LanceDBVectorStore":
kwargs["uri"] = "./.lancedb" if "uri" not in kwargs else kwargs["uri"]
kwargs["table_name"] = (
"vectors" if "table_name" not in kwargs else kwargs["table_name"]
)
vector_store = vector_store_class(**kwargs)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
return index

storage_context = StorageContext.from_defaults(vector_store=vector_store)

if documents is not None:
index = VectorStoreIndex.from_documents(
documents=documents,
storage_context=storage_context,
Expand Down
12 changes: 6 additions & 6 deletions lyzr/utils/chat_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def pdf_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -119,7 +119,7 @@ def txt_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -191,7 +191,7 @@ def docx_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -253,7 +253,7 @@ def webpage_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -315,7 +315,7 @@ def website_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -377,7 +377,7 @@ def youtube_chat_(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down
12 changes: 6 additions & 6 deletions lyzr/utils/rag_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ def pdf_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -110,7 +110,7 @@ def txt_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -175,7 +175,7 @@ def docx_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -230,7 +230,7 @@ def webpage_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -285,7 +285,7 @@ def website_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down Expand Up @@ -340,7 +340,7 @@ def youtube_rag(

llm_params = {} if llm_params is None else llm_params
vector_store_params = (
{"vector_store_type": "LanceDBVectorStore"}
{"vector_store_type": "WeaviateVectorStore"}
if vector_store_params is None
else vector_store_params
)
Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@
"llama-index==0.9.4",
"langchain==0.0.339",
"python-dotenv>=1.0.0",
"lancedb==0.3.3",
"beautifulsoup4==4.12.2",
"pandas==2.0.2",
"matplotlib==3.8.2",
"weaviate-client==3.25.3",
],
)

0 comments on commit 35440e7

Please sign in to comment.