Handle stringify_json for search output as well

OoriData · Nov 10, 2024 · b6f3bcf · b6f3bcf
1 parent c11e202
commit b6f3bcf
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/pylib/embedding/pgvector_data.py b/pylib/embedding/pgvector_data.py
@@ -235,4 +235,11 @@ async def search(
                 ),
                 *query_args
             )
+        if self.stringify_json:
+            search_results = ({
+                'cosine_similarity': r['cosine_similarity'],
+                'content': r['content'],
+                'metadata': json.loads(r['metadata']),
+                } for r in search_results)
+
         return process_search_response(search_results)
diff --git a/pylib/embedding/pgvector_message.py b/pylib/embedding/pgvector_message.py
@@ -351,7 +351,7 @@ async def get_messages(
                 'ts': record['ts'],
                 'role': record['role'],
                 'content': record['content'],
-                'metadata': record['metadata']
+                'metadata': json.loads(record['metadata']) if self.stringify_json else record['metadata']
             }) for record in message_records)
 
     async def search(
@@ -415,15 +415,15 @@ async def search(
                 *qparams
             )
 
-        search_results = [
+        search_results = (
             {
                 'ts': record['ts'],
                 'role': record['role'],
                 'content': record['content'],
-                'metadata': record['metadata'],
+                'metadata': json.loads(record['metadata']) if self.stringify_json else record['metadata'],
                 'cosine_similarity': record['cosine_similarity']
             }
             for record in message_records
-        ]
+        )
 
         return process_search_response(search_results)