Skip to content

Commit

Permalink
Merge pull request #771 from vanna-ai/revert-748-main
Browse files Browse the repository at this point in the history
Revert "feat: support pgvecto.rs"
  • Loading branch information
zainhoda authored Feb 8, 2025
2 parents 856dfa9 + f1074f7 commit 52da03d
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 278 deletions.
1 change: 0 additions & 1 deletion src/vanna/pgvector/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1 @@
from .pgvector import PG_VectorStore
from .pgvecto_rs import PG_Vecto_rsStore
269 changes: 0 additions & 269 deletions src/vanna/pgvector/pgvecto_rs.py

This file was deleted.

15 changes: 7 additions & 8 deletions src/vanna/pgvector/pgvector.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from .. import ValidationError
from ..base import VannaBase
from ..types import TrainingPlan, TrainingPlanItem
from ..utils import deterministic_uuid


class PG_VectorStore(VannaBase):
Expand Down Expand Up @@ -56,7 +55,7 @@ def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
},
ensure_ascii=False,
)
id = deterministic_uuid(question_sql_json) + "-sql"
id = str(uuid.uuid4()) + "-sql"
createdat = kwargs.get("createdat")
doc = Document(
page_content=question_sql_json,
Expand All @@ -67,7 +66,7 @@ def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
return id

def add_ddl(self, ddl: str, **kwargs) -> str:
_id = deterministic_uuid(ddl) + "-ddl"
_id = str(uuid.uuid4()) + "-ddl"
doc = Document(
page_content=ddl,
metadata={"id": _id},
Expand All @@ -76,7 +75,7 @@ def add_ddl(self, ddl: str, **kwargs) -> str:
return _id

def add_documentation(self, documentation: str, **kwargs) -> str:
_id = deterministic_uuid(documentation) + "-doc"
_id = str(uuid.uuid4()) + "-doc"
doc = Document(
page_content=documentation,
metadata={"id": _id},
Expand All @@ -95,7 +94,7 @@ def get_collection(self, collection_name):
case _:
raise ValueError("Specified collection does not exist.")

def get_similar_question_sql(self, question: str, **kwargs) -> list:
def get_similar_question_sql(self, question: str) -> list:
documents = self.sql_collection.similarity_search(query=question, k=self.n_results)
return [ast.literal_eval(document.page_content) for document in documents]

Expand Down Expand Up @@ -204,7 +203,7 @@ def remove_training_data(self, id: str, **kwargs) -> bool:
# Commit the transaction if the delete was successful
transaction.commit()
# Check if any row was deleted and return True or False accordingly
return result.rowcount() > 0
return result.rowcount > 0
except Exception as e:
# Rollback the transaction in case of error
logging.error(f"An error occurred: {e}")
Expand Down Expand Up @@ -236,9 +235,9 @@ def remove_collection(self, collection_name: str) -> bool:
try:
result = connection.execute(query)
transaction.commit() # Explicitly commit the transaction
if result.rowcount() > 0:
if result.rowcount > 0:
logging.info(
f"Deleted {result.rowcount()} rows from "
f"Deleted {result.rowcount} rows from "
f"langchain_pg_embedding where collection is {collection_name}."
)
return True
Expand Down

0 comments on commit 52da03d

Please sign in to comment.