test new extraction

lavague-ai · Aug 26, 2024 · 4ece722 · 4ece722
1 parent bf37b41
commit 4ece722
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/lavague-core/lavague/core/python_engine.py b/lavague-core/lavague/core/python_engine.py
@@ -21,6 +21,7 @@
 from llama_index.core.base.llms.base import BaseLLM
 from llama_index.core.embeddings import BaseEmbedding
 import re
+from lavague.core.extractors import JsonFromMarkdownExtractor
 
 DEFAULT_TEMPERATURE = 0.0
 
@@ -80,14 +81,7 @@ def from_context(cls, context: Context, driver: BaseDriver):
         return cls(llm=context.llm, embedding=context.embedding, driver=driver)
 
     def extract_json(self, output: str) -> Optional[dict]:
-        clean = (
-            output.replace("'ret'", '"ret"')
-            .replace("'score'", '"score"')
-            .replace("```json", "")
-            .replace("```", "")
-            .strip()
-        )
-        clean = re.sub(r"\n+", "\n", clean)
+        clean = JsonFromMarkdownExtractor.extract(output)
         try:
             output_dict = json.loads(clean)
         except json.JSONDecodeError as e: