Update QueryMetadataExtractor.py

QueryMetadataExtractor.py

@@ -11,34 +11,34 @@ class QueryMetadataExtractor:
 
     def __init__(self):
         prompt = """
-        Bạn là một phần của hệ thống thông tin xử lý các truy vấn của người dùng.
-        Với một truy vấn của người dùng, bạn trích xuất thông tin từ đó khớp với danh sách các trường siêu dữ liệu nhất định.
-        Thông tin được trích xuất từ ​​truy vấn phải khớp với ngữ nghĩa liên quan đến các trường siêu dữ liệu nhất định.
-        Thông tin bạn trích xuất từ ​​truy vấn sau đó sẽ được sử dụng làm bộ lọc để thu hẹp không gian tìm kiếm
-        khi truy vấn một chỉ mục.
-        Chỉ bao gồm giá trị của siêu dữ liệu được trích xuất mà không bao gồm tên của trường siêu dữ liệu.
-        Thông tin được trích xuất trong 'Siêu dữ liệu đã trích xuất' phải được trả về dưới dạng cấu trúc JSON hợp lệ.
-        ###
-        Ví dụ 1:
-        Truy vấn: "Doanh thu của Nvidia vào năm 2022 là bao nhiêu?"
-        Các trường siêu dữ liệu: {"company", "year"}
-        Các trường siêu dữ liệu được trích xuất: {"company": "nvidia", "year": 2022}
-        ###
-        Ví dụ 2:
-        Truy vấn: "Ấn phẩm có ảnh hưởng nhất vào năm 2023 liên quan đến bệnh Alzheimer là gì?"
-        Các trường siêu dữ liệu: {"disease", "year"}
-        Các trường siêu dữ liệu được trích xuất: {"disease": "Alzheimer's", "year": 2023}
-        ###
-        Ví dụ 3:
-        Truy vấn: "{{query}}"
-        Trường siêu dữ liệu: "{{metadata_fields}}"
-        Các trường siêu dữ liệu được trích xuất:
-        """
+          Bạn là một phần của hệ thống thông tin xử lý các truy vấn của người dùng.
+          Cho một truy vấn từ người dùng, bạn trích xuất thông tin khớp với danh sách các trường metadata đã cho.
+          Thông tin được trích xuất từ truy vấn phải khớp với ngữ nghĩa liên quan đến các trường metadata đã cho.
+          Thông tin bạn trích xuất từ truy vấn sẽ được sử dụng như bộ lọc để thu hẹp không gian tìm kiếm
+          khi truy vấn một chỉ mục.
+          Chỉ bao gồm giá trị của metadata đã trích xuất mà không bao gồm tên của trường metadata.
+          Thông tin đã trích xuất trong 'Extracted metadata' phải trả về dưới dạng cấu trúc JSON hợp lệ.
+          ###
+          Ví dụ 1:
+          Truy vấn: "Luật được Quốc hội ban hành năm 2014?"
+          Metadata fields: {"publisher", "publish_year", "document_type"}
+          Extracted metadata fields: {"publisher": "Quốc hội", "publish_year": 2014, "document_type":"Luật"}
+          ###
+          Ví dụ 2:
+          Truy vấn: "Nghị định được Chính phủ ban hành năm 2021?"
+          Trường metadata: {"publisher", "publish_year", "document_type"}
+          Extracted metadata fields: {"publisher": "Chính phủ", "publish_year": 2021, "document_type":"Nghị định"}
+          ###
+          Ví dụ 3:
+          Truy vấn: "{{query}}"
+          Trường metadata: "{{metadata_fields}}"
+          Extracted metadata fields:
+          """
         generator = OpenAIGenerator(
             api_key=Secret.from_env_var("OCTOAI_TOKEN"),
             api_base_url="https://text.octoai.run/v1",
-            model="mixtral-8x22b-instruct",
-            generation_kwargs = {"max_tokens": 512}
+            model="meta-llama-3-70b-instruct",
+            generation_kwargs = {"max_tokens": 512,"response_format": { "type": "json_object" }}
         )
         self.pipeline = Pipeline()
         self.pipeline.add_component(name="builder", instance=PromptBuilder(prompt))
@@ -48,13 +48,18 @@ class QueryMetadataExtractor:
     @component.output_types(filters=Dict[str, str])
     def run(self, query: str, metadata_fields: List[str]):
         result = self.pipeline.run({'builder': {'query': query, 'metadata_fields': metadata_fields}})
-        print(result)
-        metadata = json.loads(result['llm']['replies'][0])
-
+        
+        import json
+        first_reply = result['llm']['replies'][0]
+        parsed_data = json.loads(first_reply)
+        metadata = parsed_data['answer']
+        metadata = metadata.replace("'",'"')
+        data = json.loads(metadata)
         # this can be done with specific data structures and in a more sophisticated way
         filters = []
-        for key, value in metadata.items():
+        for key, value in data.items():
             field = f"meta.{key}"
             filters.append({f"field": field, "operator": "==", "value": value})
-
+        print(filters)
         return {"filters": {"operator": "AND", "conditions": filters}}
+