Merge pull request #82 from AbanteAI/arbitrary-role-messages

add functions to add arbitrary role to spicemessages

pyproject.toml

@@ -12,7 +12,7 @@ packages=["spice"]
 
 [project]
 name = "spiceai"
-version = "0.3.9"
+version = "0.3.10"
 license = {text = "Apache-2.0"}
 description = "A Python library for building AI-powered applications."
 readme = "README.md"

spice/spice_message.py

@@ -5,7 +5,7 @@
 from collections import UserList
 from json import JSONEncoder
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, TypedDict
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Literal, Optional, TypedDict
 
 from openai.types.chat import (
     ChatCompletionAssistantMessageParam,
@@ -25,6 +25,10 @@
 VALID_MIMETYPES = ["image/jpeg", "image/png", "image/gif", "image/webp"]
 
 
+def create_message(role: Literal["user", "assistant", "system"], content: str) -> ChatCompletionMessageParam:
+    return {"role": role, "content": content}  # pyright: ignore
+
+
 def user_message(content: str) -> ChatCompletionUserMessageParam:
     """Creates a user message with the given content."""
     return {"role": "user", "content": content}
@@ -112,6 +116,9 @@ def __init__(self, client: Spice, initlist: Optional[Iterable[SpiceMessage]] = N
         self._client = client
         super().__init__(initlist)
 
+    def add_message(self, role: Literal["user", "assistant", "system"], content: str):
+        self.data.append(create_message(role, content))
+
     def add_user_message(self, content: str):
         """Appends a user message with the given content."""
         self.data.append(user_message(content))
@@ -136,6 +143,13 @@ def add_http_image_message(self, url: str):
         """Appends a user message with the image from the given url."""
         self.data.append(http_image_message(url))
 
+    def add_prompt(self, role: Literal["user", "assistant", "system"], name: str, **context: Any):
+        prompt = self._client.get_prompt(name)
+        rendered_prompt = self._client.get_rendered_prompt(name, **context)
+        message = _MetadataDict(create_message(role, rendered_prompt))
+        message.prompt_metadata = {"name": name, "content": prompt, "context": context}
+        self.data.append(message)  # pyright: ignore
+
     def add_user_prompt(self, name: str, **context: Any):
         """Appends a user message with the given pre-loaded prompt using jinja to render the context."""
         prompt = self._client.get_prompt(name)