enhanged LLM response parsing for lists

emcf · Sep 4, 2024 · 2e52ff2 · 2e52ff2
1 parent 55a224c
commit 2e52ff2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -10,7 +10,7 @@ def read_git_requirements(file):
 
 setup(
     name='thepipe_api',
-    version='1.2.3',
+    version='1.2.5',
     author='Emmett McFarlane',
     author_email='[email protected]',
     description='AI-native extractor, powered by multimodal LLMs.',

diff --git a/thepipe/extract.py b/thepipe/extract.py
@@ -80,11 +80,15 @@ def extract_from_chunk(chunk: Chunk, chunk_index: int, schema: str, ai_model: st
                 if multiple_extractions:
                     if isinstance(llm_response_dict, dict) and "extraction" in llm_response_dict:
                         response_dict["extraction"] = llm_response_dict["extraction"]
+                    elif isinstance(llm_response_dict, list):
+                        response_dict["extraction"] = llm_response_dict
                     else:
                         response_dict["extraction"] = [llm_response_dict]
                 else:
                     if isinstance(llm_response_dict, dict):
                         response_dict.update(llm_response_dict)
+                    elif isinstance(llm_response_dict, list):
+                        response_dict["error"] = f"Expected a single JSON object but received a list: {llm_response_dict}"
                     else:
                         response_dict["error"] = f"Invalid JSON structure in LLM response: {llm_response_dict}"
             else: