ai-pipeline-core 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

ai_pipeline_core/llm/ai_messages.py

@@ -1,4 +1,5 @@
 import base64
+import hashlib
 import json
 
 from openai.types.chat import (
@@ -59,6 +60,11 @@ class AIMessages(list[AIMessageType]):
                 messages.append(message)
         return messages
 
+    def get_prompt_cache_key(self, system_prompt: str | None = None) -> str:
+        if not system_prompt:
+            system_prompt = ""
+        return hashlib.sha256((system_prompt + json.dumps(self.to_prompt())).encode()).hexdigest()
+
     @staticmethod
     def document_to_prompt(document: Document) -> list[ChatCompletionContentPartParam]:
         """

ai_pipeline_core/llm/client.py

@@ -48,15 +48,13 @@ def _process_messages(
         # Use AIMessages.to_prompt() for context
         context_messages = context.to_prompt()
 
-        # Apply caching to context messages
-        for msg in context_messages:
-            if msg.get("role") == "user":
-                # Add cache control to user messages in context
-                msg["cache_control"] = {  # type: ignore
-                    "type": "ephemeral",
-                    "ttl": "120s",  # Cache for 2m
-                }
-            processed_messages.append(msg)
+        # Apply caching to last context message
+        context_messages[-1]["cache_control"] = {  # type: ignore
+            "type": "ephemeral",
+            "ttl": "120s",  # Cache for 2m
+        }
+
+        processed_messages.extend(context_messages)
 
     # Process regular messages without caching
     if messages:
@@ -108,9 +106,14 @@ async def _generate_with_retry(
         **options.to_openai_completion_kwargs(),
     }
 
+    if context:
+        completion_kwargs["prompt_cache_key"] = context.get_prompt_cache_key(options.system_prompt)
+
     for attempt in range(options.retries):
         try:
-            with Laminar.start_as_current_span(model, span_type="LLM", input=messages) as span:
+            with Laminar.start_as_current_span(
+                model, span_type="LLM", input=processed_messages
+            ) as span:
                 response = await _generate(model, processed_messages, completion_kwargs)
                 span.set_attributes(response.get_laminar_metadata())
                 Laminar.set_span_output(response.content)

{ai_pipeline_core-0.1.3.dist-info → ai_pipeline_core-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.1.3
+Version: 0.1.4
 Summary: Core utilities for AI-powered processing pipelines using prefect
 Project-URL: Homepage, https://github.com/bbarwik/ai-pipeline-core
 Project-URL: Repository, https://github.com/bbarwik/ai-pipeline-core

{ai_pipeline_core-0.1.3.dist-info → ai_pipeline_core-0.1.4.dist-info}/RECORD

@@ -14,8 +14,8 @@ ai_pipeline_core/documents/utils.py,sha256=BdE4taSl1vrBhxnFbOP5nDA7lXIcvY__AMRTH
 ai_pipeline_core/flow/__init__.py,sha256=_Sji2yY1ICkvVX6QiiGWKzqIXtg9UAiuvhjHSK_gdO8,57
 ai_pipeline_core/flow/config.py,sha256=crbe_OvNE6qulIKv1D8yKoe8xrEsIlvICyxjhqHHBxQ,2266
 ai_pipeline_core/llm/__init__.py,sha256=3XVK-bSJdOe0s6KmmO7PDbsXHfjlcZEG1MVBmaz3EeU,442
-ai_pipeline_core/llm/ai_messages.py,sha256=lzzp4t6xDU9ULhwbmORFTQbdl0BbsqxD5YLzlHstQwU,4333
-ai_pipeline_core/llm/client.py,sha256=vVzgJXmWcWZ1limRvq_PGkFAhrfSNp18UECSiJtGhYI,7690
+ai_pipeline_core/llm/ai_messages.py,sha256=DwJJe05BtYdnMZeHbBbyEbDCqrW63SRvprxptoJUCn4,4586
+ai_pipeline_core/llm/client.py,sha256=IOcyjwyAKQWlqnwC5p2Hl4FeRCzOJAHC5Yqr_oCBQ8s,7703
 ai_pipeline_core/llm/model_options.py,sha256=TvAAlDFZN-TP9-J-RZBuU_dpSocskf6paaQMw1XY9UE,1321
 ai_pipeline_core/llm/model_response.py,sha256=fIWueaemgo0cMruvToMZyKsRPzKwL6IlvUJN7DLG710,5558
 ai_pipeline_core/llm/model_types.py,sha256=rIwY6voT8-xdfsKPDC0Gkdl2iTp9Q2LuvWGSRU9Mp3k,342
@@ -23,7 +23,7 @@ ai_pipeline_core/logging/__init__.py,sha256=DOO6ckgnMVXl29Sy7q6jhO-iW96h54pCHQDz
 ai_pipeline_core/logging/logging.yml,sha256=YTW48keO_K5bkkb-KXGM7ZuaYKiquLsjsURei8Ql0V4,1353
 ai_pipeline_core/logging/logging_config.py,sha256=6MBz9nnVNvqiLDoyy9-R3sWkn6927Re5hdz4hwTptpI,4903
 ai_pipeline_core/logging/logging_mixin.py,sha256=RDaR2ju2-vKTJRzXGa0DquGPT8_UxahWjvKJnaD0IV8,7810
-ai_pipeline_core-0.1.3.dist-info/METADATA,sha256=0TBoEep5C9LSmlOKA3koR2vfKSiIba49SLYLnVLiaWc,15869
-ai_pipeline_core-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ai_pipeline_core-0.1.3.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.1.3.dist-info/RECORD,,
+ai_pipeline_core-0.1.4.dist-info/METADATA,sha256=oB5vtkmCTKTlJKiTetHT8Lt8PKgYEAihOGIlKsD8tSQ,15869
+ai_pipeline_core-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ai_pipeline_core-0.1.4.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.1.4.dist-info/RECORD,,