morphik 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

morphik/__init__.py

@@ -2,9 +2,9 @@
 Morphik Python SDK for document ingestion and querying.
 """
 
-from .sync import Morphik
 from .async_ import AsyncMorphik
 from .models import Document
+from .sync import Morphik
 
 __all__ = [
     "Morphik",
@@ -12,4 +12,4 @@ __all__ = [
     "Document",
 ]
 
-__version__ = "0.1.4"
+__version__ = "0.1.6"

morphik/_internal.py

@@ -1,26 +1,25 @@
 import base64
 import io
 import json
-from io import BytesIO, IOBase
-from PIL import Image
-from PIL.Image import Image as PILImage
+from io import BytesIO
 from pathlib import Path
-from typing import Dict, Any, List, Optional, Union, Tuple, BinaryIO
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Type, Union
 from urllib.parse import urlparse
 
 import jwt
+from PIL import Image
+from PIL.Image import Image as PILImage
 from pydantic import BaseModel, Field
 
 from .models import (
-    Document,
     ChunkResult,
-    DocumentResult,
+    ChunkSource,  # Prompt override models
     CompletionResponse,
-    IngestTextRequest,
-    ChunkSource,
+    Document,
+    DocumentResult,
     Graph,
-    # Prompt override models
     GraphPromptOverrides,
+    IngestTextRequest,
 )
 from .rules import Rule
 
@@ -199,9 +198,7 @@ class _MorphikClientLogic:
         if rules:
             if all(isinstance(r, list) for r in rules):
                 # List of lists - per-file rules
-                converted_rules = [
-                    [self._convert_rule(r) for r in rule_list] for rule_list in rules
-                ]
+                converted_rules = [[self._convert_rule(r) for r in rule_list] for rule_list in rules]
             else:
                 # Flat list - shared rules for all files
                 converted_rules = [self._convert_rule(r) for r in rules]
@@ -235,8 +232,9 @@ class _MorphikClientLogic:
         hop_depth: int,
         include_paths: bool,
         prompt_overrides: Optional[Dict],
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
+        schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
     ) -> Dict[str, Any]:
         """Prepare request for query endpoint"""
         payload = {
@@ -256,6 +254,20 @@ class _MorphikClientLogic:
             payload["folder_name"] = folder_name
         if end_user_id:
             payload["end_user_id"] = end_user_id
+
+        # Add schema to payload if provided
+        if schema:
+            # If schema is a Pydantic model class, serialize it to a JSON schema dict
+            if isinstance(schema, type) and issubclass(schema, BaseModel):
+                payload["schema"] = schema.model_json_schema()
+            elif isinstance(schema, dict):
+                # Basic check if it looks like a JSON schema (has 'properties' or 'type')
+                if "properties" not in schema and "type" not in schema:
+                    raise ValueError("Provided schema dictionary does not look like a valid JSON schema")
+                payload["schema"] = schema
+            else:
+                raise TypeError("schema must be a Pydantic model type or a dictionary representing a JSON schema")
+
         # Filter out None values before sending
         return {k_p: v_p for k_p, v_p in payload.items() if v_p is not None}
 
@@ -266,7 +278,7 @@ class _MorphikClientLogic:
         k: int,
         min_score: float,
         use_colpali: bool,
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Dict[str, Any]:
         """Prepare request for retrieve_chunks endpoint"""
@@ -290,7 +302,7 @@ class _MorphikClientLogic:
         k: int,
         min_score: float,
         use_colpali: bool,
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Dict[str, Any]:
         """Prepare request for retrieve_docs endpoint"""
@@ -312,7 +324,7 @@ class _MorphikClientLogic:
         skip: int,
         limit: int,
         filters: Optional[Dict[str, Any]],
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         """Prepare request for list_documents endpoint"""
@@ -328,7 +340,7 @@ class _MorphikClientLogic:
         return params, data
 
     def _prepare_batch_get_documents_request(
-        self, document_ids: List[str], folder_name: Optional[str], end_user_id: Optional[str]
+        self, document_ids: List[str], folder_name: Optional[Union[str, List[str]]], end_user_id: Optional[str]
     ) -> Dict[str, Any]:
         """Prepare request for batch_get_documents endpoint"""
         if folder_name or end_user_id:
@@ -343,7 +355,7 @@ class _MorphikClientLogic:
     def _prepare_batch_get_chunks_request(
         self,
         sources: List[Union[ChunkSource, Dict[str, Any]]],
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Dict[str, Any]:
         """Prepare request for batch_get_chunks endpoint"""
@@ -361,7 +373,8 @@ class _MorphikClientLogic:
             if end_user_id:
                 request["end_user_id"] = end_user_id
             return request
-        return source_dicts  # Return just sources list if no scoping is needed
+        # Return the dictionary structure { "sources": [...] } consistently.
+        return {"sources": source_dicts}
 
     def _prepare_create_graph_request(
         self,
@@ -369,7 +382,7 @@ class _MorphikClientLogic:
         filters: Optional[Dict[str, Any]],
         documents: Optional[List[str]],
         prompt_overrides: Optional[Union[GraphPromptOverrides, Dict[str, Any]]],
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Dict[str, Any]:
         """Prepare request for create_graph endpoint"""
@@ -395,7 +408,7 @@ class _MorphikClientLogic:
         additional_filters: Optional[Dict[str, Any]],
         additional_documents: Optional[List[str]],
         prompt_overrides: Optional[Union[GraphPromptOverrides, Dict[str, Any]]],
-        folder_name: Optional[str],
+        folder_name: Optional[Union[str, List[str]]],
         end_user_id: Optional[str],
     ) -> Dict[str, Any]:
         """Prepare request for update_graph endpoint"""
@@ -454,15 +467,11 @@ class _MorphikClientLogic:
         docs = [Document(**doc) for doc in response_json]
         return docs
 
-    def _parse_document_result_list_response(
-        self, response_json: List[Dict[str, Any]]
-    ) -> List[DocumentResult]:
+    def _parse_document_result_list_response(self, response_json: List[Dict[str, Any]]) -> List[DocumentResult]:
         """Parse document result list response"""
         return [DocumentResult(**r) for r in response_json]
 
-    def _parse_chunk_result_list_response(
-        self, response_json: List[Dict[str, Any]]
-    ) -> List[FinalChunkResult]:
+    def _parse_chunk_result_list_response(self, response_json: List[Dict[str, Any]]) -> List[FinalChunkResult]:
         """Parse chunk result list response"""
         chunks = [ChunkResult(**r) for r in response_json]