add sentence-transformer

babyagi/classesa/diamond.py

@@ -31,7 +31,7 @@ class ProductDatabase:
     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))
+            cursor.execute("UPDATE diamondprice SET vector_col = %s WHERE id = %s", (vector, product_id))
             self.conn.commit()
 
     def search_similar_vectors(self, query_text, top_k=5):
@@ -39,13 +39,23 @@ class ProductDatabase:
         with self.conn.cursor() as cursor:
             cursor.execute("""
                 SELECT id, vector_col <=> %s::vector AS distance
-                FROM products
+                FROM diamondprice
                 ORDER BY distance
                 LIMIT %s;
             """, (query_vector, top_k))
             results = cursor.fetchall()
             return results
 
+    def search_similar_all(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,'carat', 'cut', 'color', 'clarity', 'depth', 'diamondprice.table', 'x', 'y', 'z'
+                FROM diamondprice
+            """, (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"
@@ -60,12 +70,14 @@ def main():
         # pgvector拡張機能のインストールとカラムの追加
         db.setup_vector_extension_and_column()
         print("Vector extension installed and column added successfully.")
-        
+        query_text="1"
+        results = db.search_similar_all(query_text)
+        print("Search results:")
+        for result in results:
+            print(result)        
+        return
         # サンプルデータの挿入
-        sample_text = """検査にはどのぐらい時間かかりますか？⇒当日に分かります。
-法人取引やってますか？⇒大丈夫ですよ。成約時に必要な書類の説明
-LINEで金粉送って、査定はできますか？⇒できますが、今お話した内容と同様で、検査が必要な旨を返すだけなので、金粉ではなく、他のお品物でLINE査定くださいと。
-分かりました、またどうするか検討して連絡しますと"""
+        sample_text = """"""
         sample_product_id = 1  # 実際の製品IDを使用
         db.insert_vector(sample_product_id, sample_text)
         db.insert_vector(2, sample_text)

requirements.txt

@@ -56,4 +56,5 @@ torchvision
 transformers
 langchain
 langchain_groq
-sqlalchemy
+sqlalchemy
+sentence-transformers