cleant: split article model

custom_pgvector.py

@@ -9,21 +9,16 @@ import pandas as pd
 import sqlalchemy
 from langchain.docstore.document import Document
 from langchain.schema.embeddings import Embeddings
-from langchain.utils import get_from_dict_or_env
 from langchain.vectorstores.base import VectorStore
-from pgvector.sqlalchemy import Vector
 from sqlalchemy import delete, text
-from sqlalchemy.orm import Session, declarative_base
+from sqlalchemy.orm import Session
 
 from utils import str_to_list
-from distance import DistanceStrategy
-
+from models.distance import DistanceStrategy
+from model import Base, Article
 
 DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.COSINE
 
-Base = declarative_base()  # type: Any
-
-
 _LANGCHAIN_DEFAULT_COLLECTION_NAME = "langchain"
 
 
@@ -32,18 +27,6 @@ def _results_to_docs(docs_and_scores: Any) -> List[Document]:
     return [doc for doc, _ in docs_and_scores]
 
 
-class Article(Base):
-    """Embedding store."""
-
-    __tablename__ = "article"
-
-    id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True, nullable=False)
-    title = sqlalchemy.Column(sqlalchemy.String, nullable=True)
-    abstract = sqlalchemy.Column(sqlalchemy.String, nullable=True)
-    embedding: Vector = sqlalchemy.Column("abstract_embedding", Vector(None))
-    doi = sqlalchemy.Column(sqlalchemy.String, nullable=True)
-
-
 class CustomPGVector(VectorStore):
     """`Postgres`/`PGVector` vector store.
 

model.py

@@ -0,0 +1,18 @@
+import sqlalchemy
+
+from sqlalchemy.orm import declarative_base
+from pgvector.sqlalchemy import Vector
+
+Base = declarative_base()  # type: Any
+
+
+class Article(Base):
+    """Embedding store."""
+
+    __tablename__ = "article"
+
+    id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True, nullable=False)
+    title = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+    abstract = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+    embedding: Vector = sqlalchemy.Column("abstract_embedding", Vector(None))
+    doi = sqlalchemy.Column(sqlalchemy.String, nullable=True)