fda

babyagi/classesa/diamond.py

@@ -0,0 +1,88 @@
+import psycopg2
+from sentence_transformers import SentenceTransformer
+
+class ProductDatabase:
+    def __init__(self, database_url):
+        self.database_url = database_url
+        self.conn = None
+        self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    
+    def connect(self):
+        self.conn = psycopg2.connect(self.database_url)
+    
+    def close(self):
+        if self.conn:
+            self.conn.close()
+    
+    def setup_vector_extension_and_column(self):
+        with self.conn.cursor() as cursor:
+            # pgvector拡張機能のインストール
+            cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+            
+            # ベクトルカラムの追加
+            cursor.execute("ALTER TABLE products ADD COLUMN IF NOT EXISTS vector_col vector(384);")
+            
+            self.conn.commit()
+
+    def get_embedding(self, text):
+        embedding = self.model.encode(text)
+        return embedding
+
+    def insert_vector(self, product_id, text):
+        vector = self.get_embedding(text).tolist()  # ndarray をリストに変換
+        with self.conn.cursor() as cursor:
+            cursor.execute("UPDATE products SET vector_col = %s WHERE id = %s", (vector, product_id))
+            self.conn.commit()
+
+    def search_similar_vectors(self, query_text, top_k=5):
+        query_vector = self.get_embedding(query_text).tolist()  # ndarray をリストに変換
+        with self.conn.cursor() as cursor:
+            cursor.execute("""
+                SELECT id, vector_col <=> %s::vector AS distance
+                FROM products
+                ORDER BY distance
+                LIMIT %s;
+            """, (query_vector, top_k))
+            results = cursor.fetchall()
+            return results
+
+def main():
+    # データベース接続情報
+    DATABASE_URL = "postgresql://miyataken999:yz1wPf4KrWTm@ep-odd-mode-93794521.us-east-2.aws.neon.tech/neondb?sslmode=require"
+    
+    # ProductDatabaseクラスのインスタンスを作成
+    db = ProductDatabase(DATABASE_URL)
+    
+    # データベースに接続
+    db.connect()
+    
+    try:
+        # pgvector拡張機能のインストールとカラムの追加
+        db.setup_vector_extension_and_column()
+        print("Vector extension installed and column added successfully.")
+        
+        # サンプルデータの挿入
+        sample_text = """検査にはどのぐらい時間かかりますか？⇒当日に分かります。
+法人取引やってますか？⇒大丈夫ですよ。成約時に必要な書類の説明
+LINEで金粉送って、査定はできますか？⇒できますが、今お話した内容と同様で、検査が必要な旨を返すだけなので、金粉ではなく、他のお品物でLINE査定くださいと。
+分かりました、またどうするか検討して連絡しますと"""
+        sample_product_id = 1  # 実際の製品IDを使用
+        db.insert_vector(sample_product_id, sample_text)
+        db.insert_vector(2, sample_text)
+
+        print(f"Vector inserted for product ID {sample_product_id}.")
+
+        
+        # ベクトル検索
+        query_text = "今お話した内容と同様で"
+        results = db.search_similar_vectors(query_text)
+        print("Search results:")
+        for result in results:
+            print(result)
+    
+    finally:
+        # 接続を閉じる
+        db.close()
+
+if __name__ == "__main__":
+    main()