test

babyagi/babyagi.py

@@ -24,6 +24,31 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 import numpy
 
+import psycopg2
+
+class ProductDatabase:
+    def __init__(self, database_url):
+        self.database_url = database_url
+        self.conn = None
+    
+    def connect(self):
+        self.conn = psycopg2.connect(self.database_url)
+    
+    def close(self):
+        if self.conn:
+            self.conn.close()
+    
+    def fetch_data(self):
+        with self.conn.cursor() as cursor:
+            cursor.execute("SELECT id FROM products")
+            rows = cursor.fetchall()
+            return rows
+    
+    def update_data(self, product_id, new_price):
+        with self.conn.cursor() as cursor:
+            cursor.execute("UPDATE products SET price = %s WHERE id = %s", (new_price, product_id))
+            self.conn.commit()
+
 # モデル名を指定
 model_name = "sentence-transformers/all-MiniLM-L6-v2"
 
@@ -745,3 +770,32 @@ def main():
 
 if __name__ == "__main__":
     main()
+
+def test_postgres():
+    # データベース接続情報
+    DATABASE_URL = "postgresql://miyataken999:[email protected]/neondb?sslmode=require"
+    
+    # ProductDatabaseクラスのインスタンスを作成
+    db = ProductDatabase(DATABASE_URL)
+    
+    # データベースに接続
+    db.connect()
+    
+    try:
+        # データを取得
+        products = db.fetch_data()
+        print("Fetched products:")
+        for product in products:
+            print(product)
+        
+        # データを更新（例: 価格を更新）
+        for product in products:
+            product_id = product[0]
+            print(product_id)
+            #new_price = product[2] * 1.1  # 価格を10%増加させる
+            #db.update_data(product_id, new_price)
+            #print(f"Updated product ID {product_id} with new price {new_price}")
+    
+    finally:
+        # 接続を閉じる
+        db.close()

babyagi/classesa/vector.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:[email protected]/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()