pixie-prompts 0.1.10.dev1__tar.gz → 0.1.12__tar.gz

{pixie_prompts-0.1.10.dev1 → pixie_prompts-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pixie-prompts
-Version: 0.1.10.dev1
+Version: 0.1.12
 Summary: Code-first, type-safe prompt management
 License: MIT
 License-File: LICENSE

{pixie_prompts-0.1.10.dev1 → pixie_prompts-0.1.12}/pixie/prompts/graphql.py

@@ -95,6 +95,22 @@ class LlmCallResult:
     reasoning: str | None
 
 
+@strawberry.input
+class LlmCallInput:
+    """Input for a single LLM call in a batch."""
+
+    id: strawberry.ID
+    """Unique identifier for this call, used to correlate results."""
+    model: str
+    prompt_template: str
+    variables: Optional[JSON] = None
+    prompt_placeholder: str = "{{prompt}}"
+    input_messages: list[JSON] = strawberry.field(default_factory=list)
+    output_schema: Optional[JSON] = None
+    tools: Optional[list[JSON]] = None
+    model_parameters: Optional[JSON] = None
+
+
 def is_demo_mode() -> bool:
     is_demo_mode = os.getenv("IS_DEMO_MODE", "0") in ("1", "true", "True")
     return is_demo_mode
@@ -227,77 +243,20 @@ class Mutation:
             GraphQLError: If the LLM call fails.
         """
         try:
-            if is_demo_mode():
-                model = "openai:gpt-4o-mini"
-            template = jinja2.Template(prompt_template)
-            prompt = template.render(**(cast(dict[str, Any], variables) or {}))
-            print(prompt)
-            print(type(prompt))
-            pydantic_messages = openai_messages_to_pydantic_ai_messages(
-                cast(list[dict[str, Any]], input_messages)
-            )
-            for msg in pydantic_messages:
-                for part in msg.parts:
-                    if part.part_kind == "user-prompt":
-                        if isinstance(part.content, str):
-                            part.content = part.content.replace(
-                                prompt_placeholder,
-                                prompt,
-                            )
-                        else:
-                            part.content = [
-                                p.replace(prompt_placeholder, prompt)
-                                for p in part.content
-                                if isinstance(p, str)
-                            ]
-                    elif part.part_kind == "system-prompt":
-                        part.content = part.content.replace(prompt_placeholder, prompt)
-
-            # Replace the placeholder in input messages
-            response = await model_request(
+            call_input = LlmCallInput(
+                id=strawberry.ID("single-call"),
                 model=model,
-                messages=pydantic_messages,
-                model_settings=cast(ModelSettings | None, model_parameters),
-                model_request_parameters=assemble_model_request_parameters(
-                    cast(dict[str, Any] | None, output_schema),
-                    cast(list[dict[str, Any]] | None, tools),
-                    strict=True,
-                    allow_text_output=False,
-                ),
-            )
-            return LlmCallResult(
-                input=JSON(pydantic_ai_messages_to_openai_messages(pydantic_messages)),
-                output=(
-                    JSON(json.loads(response.text) if output_schema else response.text)
-                    if response.text
-                    else None
-                ),
-                tool_calls=(
-                    [
-                        ToolCall(
-                            name=tc.tool_name,
-                            args=JSON(tc.args_as_dict()),
-                            tool_call_id=strawberry.ID(tc.tool_call_id),
-                        )
-                        for tc in response.tool_calls
-                    ]
-                    if response.tool_calls
-                    else None
-                ),
-                usage=JSON(
-                    {
-                        "input_tokens": response.usage.input_tokens,
-                        "output_tokens": response.usage.output_tokens,
-                        "total_tokens": response.usage.total_tokens,
-                    }
-                ),
-                cost=float(response.cost().total_price),
-                timestamp=response.timestamp,
-                reasoning=response.thinking,
+                prompt_template=prompt_template,
+                variables=variables,
+                prompt_placeholder=prompt_placeholder,
+                input_messages=input_messages,
+                output_schema=output_schema,
+                tools=tools,
+                model_parameters=model_parameters,
             )
+            return await execute_single_llm_call(call_input)
         except Exception as e:
-            logger.error("Error running LLM: %s", str(e))
-            raise GraphQLError(f"Failed to run LLM: {str(e)}") from e
+            raise GraphQLError(f"LLM call failed: {str(e)}") from e
 
     @strawberry.mutation
     async def add_prompt_version(
@@ -364,5 +323,85 @@ class Mutation:
         return "OK"
 
 
+async def execute_single_llm_call(
+    call_input: LlmCallInput,
+) -> LlmCallResult:
+    """Execute a single LLM call and return the result as a BatchLlmCallUpdate."""
+    model = call_input.model
+    template = jinja2.Template(call_input.prompt_template)
+    prompt = template.render(**(cast(dict[str, Any], call_input.variables) or {}))
+
+    pydantic_messages = openai_messages_to_pydantic_ai_messages(
+        cast(list[dict[str, Any]], call_input.input_messages)
+    )
+
+    # Replace the placeholder in messages
+    for msg in pydantic_messages:
+        for part in msg.parts:
+            if part.part_kind == "user-prompt":
+                if isinstance(part.content, str):
+                    part.content = part.content.replace(
+                        call_input.prompt_placeholder,
+                        prompt,
+                    )
+                else:
+                    part.content = [
+                        p.replace(call_input.prompt_placeholder, prompt)
+                        for p in part.content
+                        if isinstance(p, str)
+                    ]
+            elif part.part_kind == "system-prompt":
+                part.content = part.content.replace(
+                    call_input.prompt_placeholder, prompt
+                )
+
+    response = await model_request(
+        model=model,
+        messages=pydantic_messages,
+        model_settings=cast(ModelSettings | None, call_input.model_parameters),
+        model_request_parameters=assemble_model_request_parameters(
+            cast(dict[str, Any] | None, call_input.output_schema),
+            cast(list[dict[str, Any]] | None, call_input.tools),
+            strict=True,
+            allow_text_output=False,
+        ),
+    )
+
+    result = LlmCallResult(
+        input=JSON(pydantic_ai_messages_to_openai_messages(pydantic_messages)),
+        output=(
+            JSON(
+                json.loads(response.text) if call_input.output_schema else response.text
+            )
+            if response.text
+            else None
+        ),
+        tool_calls=(
+            [
+                ToolCall(
+                    name=tc.tool_name,
+                    args=JSON(tc.args_as_dict()),
+                    tool_call_id=strawberry.ID(tc.tool_call_id),
+                )
+                for tc in response.tool_calls
+            ]
+            if response.tool_calls
+            else None
+        ),
+        usage=JSON(
+            {
+                "input_tokens": response.usage.input_tokens,
+                "output_tokens": response.usage.output_tokens,
+                "total_tokens": response.usage.total_tokens,
+            }
+        ),
+        cost=float(response.cost().total_price),
+        timestamp=response.timestamp,
+        reasoning=response.thinking,
+    )
+
+    return result
+
+
 # Create the schema
 schema = strawberry.Schema(query=Query, mutation=Mutation)

{pixie_prompts-0.1.10.dev1 → pixie_prompts-0.1.12}/pixie/prompts/prompt.py

@@ -49,7 +49,7 @@ key is the id() of the compiled string."""
 def _find_matching_prompt(obj):
     if isinstance(obj, str):
         for compiled in _compiled_prompt_registry.values():
-            if compiled.value == obj:
+            if compiled.value in obj:
                 return compiled
         return None
     elif isinstance(obj, dict):
@@ -76,7 +76,7 @@ def get_compiled_prompt(text: str) -> CompiledPrompt | None:
     if direct_match:
         return direct_match
     for compiled in _compiled_prompt_registry.values():
-        if compiled.value == text:
+        if compiled.value in text:
             return compiled
     try:
         obj = json.loads(text)

{pixie_prompts-0.1.10.dev1 → pixie_prompts-0.1.12}/pixie/prompts/storage.py

@@ -41,6 +41,10 @@ class PromptLoadError(Exception):
         super().__init__(message)
 
 
+class PromptNotFoundError(KeyError):
+    pass
+
+
 class BaseUntypedPromptWithCreationTime(BaseUntypedPrompt):
 
     def __init__(
@@ -254,7 +258,12 @@ class _FilePromptStorage(PromptStorage):
         return original is None
 
     def get(self, prompt_id: str) -> BaseUntypedPromptWithCreationTime:
-        return self._prompts[prompt_id]
+        try:
+            return self._prompts[prompt_id]
+        except KeyError:
+            raise PromptNotFoundError(
+                f"Prompt with ID '{prompt_id}' not found in storage."
+            )
 
 
 _storage_instance: PromptStorage | None = None
@@ -306,9 +315,11 @@ class StorageBackedPrompt(Prompt[TPromptVar]):
         id: str,
         *,
         variables_definition: type[TPromptVar] = NoneType,
+        default: str | None = None,
     ) -> None:
         self._id = id
         self._variables_definition = variables_definition
+        self._default = default
         self._prompt: BasePrompt[TPromptVar] | None = None
 
     @property
@@ -319,24 +330,42 @@ class StorageBackedPrompt(Prompt[TPromptVar]):
     def variables_definition(self) -> type[TPromptVar]:
         return self._variables_definition
 
+    @property
+    def default(self) -> str | None:
+        return self._default
+
     def get_variables_schema(self) -> dict[str, Any]:
         return variables_definition_to_schema(self._variables_definition)
 
     def _get_prompt(self) -> BasePrompt[TPromptVar]:
-        storage = _ensure_storage_initialized()
-        if self._prompt is None:
-            untyped_prompt = storage.get(self.id)
-            self._prompt = BasePrompt.from_untyped(
-                untyped_prompt,
+        if self._prompt is not None:
+            return self._prompt
+        try:
+            storage = _ensure_storage_initialized()
+            if self._prompt is None:
+                untyped_prompt = storage.get(self.id)
+                self._prompt = BasePrompt.from_untyped(
+                    untyped_prompt,
+                    variables_definition=self.variables_definition,
+                )
+                schema_from_storage = untyped_prompt.get_variables_schema()
+                schema_from_definition = self.get_variables_schema()
+                if not isSubschema(schema_from_definition, schema_from_storage):
+                    raise TypeError(
+                        "Schema from definition is not a subschema of the schema from storage."
+                    )
+            return self._prompt
+        except (PromptNotFoundError, PromptLoadError):
+            if self.default is None:
+                raise PromptNotFoundError(
+                    f"Cannot load prompt with id '{self.id}' from storage, and no default is provided."
+                )
+            self._prompt = BasePrompt(
+                id=self.id,
+                versions={"v0": self.default},
                 variables_definition=self.variables_definition,
             )
-            schema_from_storage = untyped_prompt.get_variables_schema()
-            schema_from_definition = self.get_variables_schema()
-            if not isSubschema(schema_from_definition, schema_from_storage):
-                raise TypeError(
-                    "Schema from definition is not a subschema of the schema from storage."
-                )
-        return self._prompt
+            return self._prompt
 
     def actualize(self) -> Self:
         self._get_prompt()

{pixie_prompts-0.1.10.dev1 → pixie_prompts-0.1.12}/pyproject.toml

@@ -4,7 +4,7 @@ packages = [
     { include = "pixie" },
 ]
 
-version = "0.1.10.dev1"
+version = "0.1.12"
 description = "Code-first, type-safe prompt management"
 authors = ["Yiou Li <yol@gopixie.ai>"]
 license = "MIT"