fixed a lot of bugs and created comprehensive example

Author: olliekm

examples/comprehensive_demo.py

@@ -32,6 +32,7 @@ dependencies = [
     "aiohttp>=3.9.0",
     "pyyaml>=6.0.0",
     "scipy>=1.10.0",
+    "python-dotenv (>=1.2.1,<2.0.0)",
 ]
 
 [project.urls]

poetry.lock

@@ -2,6 +2,7 @@
 from typing import Any, Optional
 import hashlib
 import json
+from pydantic import BaseModel
 
 
 def generate_cache_key(
@@ -44,10 +45,20 @@ def generate_cache_key(
         - All components are sorted to ensure deterministic ordering
     """
     normalized_prompt = prompt.strip()
+
+    # Handle Pydantic models by converting to JSON schema
+    if schema:
+        if isinstance(schema, type) and issubclass(schema, BaseModel):
+            schema_str = json.dumps(schema.model_json_schema(), sort_keys=True)
+        else:
+            schema_str = json.dumps(schema, sort_keys=True)
+    else:
+        schema_str = ""
+
     key_components = {
         "prompt": normalized_prompt,
         "model": model,
-        "schema": json.dumps(schema, sort_keys=True) if schema else "",
+        "schema": schema_str,
         "temperature": temperature,
         **kwargs
     }

pyproject.toml

@@ -107,9 +107,14 @@ async def _generate():
 
                 if validation.status == ValidationStatus.VALID:
                     if self.collector:
+                        # Convert Pydantic model to JSON schema for storage
+                        schema_for_storage = schema
+                        if isinstance(schema, type) and issubclass(schema, BaseModel):
+                            schema_for_storage = schema.model_json_schema()
+
                         self.collector.collect({
                             "prompt": prompt,
-                            "json_schema": schema,
+                            "json_schema": schema_for_storage,
                             "response": generation.output,
                             "parsed_output": validation.parsed_output,
                             "success": True,
@@ -153,4 +158,18 @@ async def _generate():
 
                 # Will retry on next iteration
                 retry_count += 1
-                continue
+                continue
+
+        # If we exit the loop without returning, all retries were exhausted
+        # Return the last failed result
+        if last_validation:
+            return EnforcedOutput(
+                data=last_validation.parsed_output,
+                generation=generation,
+                validation=last_validation,
+                retry_count=retry_count,
+                success=False
+            )
+        else:
+            # Should never reach here, but handle it gracefully
+            raise RuntimeError("Enforcement failed: no validation results available")

src/parsec/cache/keys.py

@@ -2,7 +2,9 @@
 from parsec.core import BaseLLMAdapter, GenerationResponse, ModelProviders
 from typing import AsyncIterator
 from parsec.logging import get_logger
+from pydantic import BaseModel
 import time
+import json
 
 class AnthropicAdapter(BaseLLMAdapter):
     """Adapter for Anthropic's API with custom configurations."""
@@ -47,8 +49,12 @@ async def generate(self, prompt: str, schema=None, temperature=0.7,
             # Anthropic doesn't have response_format parameter
             # Instead, we need to instruct it via the prompt
             if schema:
-                import json
-                schema_str = json.dumps(schema, indent=2)
+                # Convert Pydantic model to JSON schema if needed
+                if isinstance(schema, type) and issubclass(schema, BaseModel):
+                    schema_dict = schema.model_json_schema()
+                else:
+                    schema_dict = schema
+                schema_str = json.dumps(schema_dict, indent=2)
                 message_params["messages"][0]["content"] = (
                     f"{prompt}\n\n"
                     f"Please respond with valid JSON matching this schema:\n{schema_str}\n"

src/parsec/enforcement/engine.py

@@ -2,6 +2,7 @@
 from parsec.core import BaseLLMAdapter, GenerationResponse, ModelProviders
 from typing import AsyncIterator
 from parsec.logging import get_logger
+from pydantic import BaseModel
 import time
 import json
 
@@ -69,7 +70,12 @@ async def generate(
         # If schema is provided, use JSON mode and add schema to prompt
         if schema:
             generation_config["response_mime_type"] = "application/json"
-            schema_str = json.dumps(schema, indent=2)
+            # Convert Pydantic model to JSON schema if needed
+            if isinstance(schema, type) and issubclass(schema, BaseModel):
+                schema_dict = schema.model_json_schema()
+            else:
+                schema_dict = schema
+            schema_str = json.dumps(schema_dict, indent=2)
             prompt = (
                 f"{prompt}\n\n"
                 f"Please respond with valid JSON matching this schema:\n{schema_str}\n"
@@ -140,7 +146,12 @@ async def generate_stream(
         # If schema is provided, use JSON mode and add schema to prompt
         if schema:
             generation_config["response_mime_type"] = "application/json"
-            schema_str = json.dumps(schema, indent=2)
+            # Convert Pydantic model to JSON schema if needed
+            if isinstance(schema, type) and issubclass(schema, BaseModel):
+                schema_dict = schema.model_json_schema()
+            else:
+                schema_dict = schema
+            schema_str = json.dumps(schema_dict, indent=2)
             prompt = (
                 f"{prompt}\n\n"
                 f"Please respond with valid JSON matching this schema:\n{schema_str}\n"

src/parsec/models/adapters/anthropic_adapter.py

@@ -2,6 +2,7 @@
 from parsec.core import BaseLLMAdapter, GenerationResponse, ModelProviders
 from typing import AsyncIterator
 from parsec.logging import get_logger
+from pydantic import BaseModel
 import time
 import json
 
@@ -41,8 +42,13 @@ async def generate(self, prompt: str, schema=None, temperature=0.7,
         extra_args = {}
         if schema and self.supports_native_structure_output():
             extra_args["response_format"] = {"type": "json_object"}
+            # Convert Pydantic model to JSON schema if needed
+            if isinstance(schema, type) and issubclass(schema, BaseModel):
+                schema_dict = schema.model_json_schema()
+            else:
+                schema_dict = schema
             # Add schema to prompt
-            messages[0]["content"] = f"{prompt}\n\nReturn valid JSON matching this schema: {json.dumps(schema)}"
+            messages[0]["content"] = f"{prompt}\n\nReturn valid JSON matching this schema: {json.dumps(schema_dict)}"
         try:
             response = await client.chat.completions.create(
                 model=self.model,
@@ -85,7 +91,12 @@ async def generate_stream(
         extra_args = {}
         if schema and self.supports_native_structure_output():
             extra_args["response_format"] = {"type": "json_object"}
-            messages[0]["content"] = f"{prompt}\n\nReturn valid JSON matching this schema: {json.dumps(schema)}"
+            # Convert Pydantic model to JSON schema if needed
+            if isinstance(schema, type) and issubclass(schema, BaseModel):
+                schema_dict = schema.model_json_schema()
+            else:
+                schema_dict = schema
+            messages[0]["content"] = f"{prompt}\n\nReturn valid JSON matching this schema: {json.dumps(schema_dict)}"
 
         stream = await client.chat.completions.create(
             model=self.model,