diff --git a/README.md b/README.md index d7f8233..5f447a5 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ You can override the shared client by calling the init function. import workflowai workflowai.init( - url=..., # defaults to WORKFLOWAI_API_URL env var or https://api.workflowai.com + url=..., # defaults to WORKFLOWAI_API_URL env var or https://run.workflowai.com (our [globally distributed, highly available endpoint](https://docs.workflowai.com/workflowai-cloud/reliability)) api_key=..., # defaults to WORKFLOWAI_API_KEY env var ) ``` @@ -335,7 +335,7 @@ image = Image(content_type='image/jpeg', data='') image = Image(url="https://example.com/image.jpg") ``` -An example of using image as input is available in [city_identifier.py](./examples/images/city_identifier.py). +An example of using image as input is available in [07_image_agent.py](./examples/07_image_agent.py). ### Files (PDF, .txt, ...) diff --git a/examples/01_basic_agent.py b/examples/01_basic_agent.py index 239ed4b..93794cc 100644 --- a/examples/01_basic_agent.py +++ b/examples/01_basic_agent.py @@ -61,13 +61,13 @@ async def main(): # Example 1: Basic usage with Paris print("\nExample 1: Basic usage with Paris") print("-" * 50) - run = await get_capital_info(CityInput(city="Paris")) + run = await get_capital_info.run(CityInput(city="Paris")) print(run) # Example 2: Using Tokyo print("\nExample 2: Using Tokyo") print("-" * 50) - run = await get_capital_info(CityInput(city="Tokyo")) + run = await get_capital_info.run(CityInput(city="Tokyo")) print(run) diff --git a/examples/07_image_agent.py b/examples/07_image_agent.py index 6a324ee..b6e5683 100644 --- a/examples/07_image_agent.py +++ b/examples/07_image_agent.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field # pyright: ignore [reportUnknownVariableType] import workflowai -from workflowai import Run, WorkflowAIError +from workflowai import WorkflowAIError from workflowai.core.domain.model import Model from workflowai.fields import Image @@ -31,8 +31,8 @@ class ImageOutput(BaseModel): ) -@workflowai.agent(id="city-identifier", model=Model.GEMINI_1_5_FLASH_LATEST) -async def identify_city_from_image(_: ImageInput) -> Run[ImageOutput]: +@workflowai.agent(id="city-identifier", model=Model.GEMINI_2_0_FLASH_LATEST) +async def identify_city_from_image(image_input: ImageInput) -> ImageOutput: """ Analyze the provided image and identify the city and country shown in it. If the image shows a recognizable landmark or cityscape, identify the city and country. @@ -62,9 +62,8 @@ async def main(): image = Image(content_type="image/jpeg", data=content) try: - agent_run = await identify_city_from_image( + agent_run = await identify_city_from_image.run( ImageInput(image=image), - use_cache="auto", ) except WorkflowAIError as e: print(f"Failed to run task. Code: {e.error.code}. Message: {e.error.message}") @@ -77,9 +76,8 @@ async def main(): # Example using URL for Image image_url = "https://t4.ftcdn.net/jpg/02/96/15/35/360_F_296153501_B34baBHDkFXbl5RmzxpiOumF4LHGCvAE.jpg" image = Image(url=image_url) - agent_run = await identify_city_from_image( + agent_run = await identify_city_from_image.run( ImageInput(image=image), - use_cache="auto", ) print("\n--------\nAgent output:\n", agent_run.output, "\n--------\n") diff --git a/examples/11_ecommerce_chatbot.py b/examples/11_ecommerce_chatbot.py index e67685f..d41f751 100644 --- a/examples/11_ecommerce_chatbot.py +++ b/examples/11_ecommerce_chatbot.py @@ -2,25 +2,17 @@ This example demonstrates how to create an e-commerce chatbot that: 1. Understands customer queries about products 2. Provides helpful responses with product recommendations -3. Maintains context through conversation +3. Maintains context through conversation using .reply 4. Returns structured product recommendations """ import asyncio -from enum import Enum from typing import Optional from pydantic import BaseModel, Field import workflowai -from workflowai import Model, Run - - -class Role(str, Enum): - """Enum representing possible message roles.""" - - USER = "user" - ASSISTANT = "assistant" +from workflowai import Model class Product(BaseModel): @@ -56,16 +48,11 @@ class Product(BaseModel): ) -class Message(BaseModel): - """Model representing a chat message.""" +class AssistantMessage(BaseModel): + """Model representing a message from the assistant.""" - role: Role = Field() content: str = Field( description="The content of the message", - examples=[ - "I'm looking for noise-cancelling headphones for travel", - "Based on your requirements, here are some great headphone options...", - ], ) recommended_products: Optional[list[Product]] = Field( default=None, @@ -73,13 +60,6 @@ class Message(BaseModel): ) -class AssistantMessage(Message): - """Model representing a message from the assistant.""" - - role: Role = Role.ASSISTANT - content: str = "" - - class ChatbotOutput(BaseModel): """Output model for the chatbot response.""" @@ -89,12 +69,8 @@ class ChatbotOutput(BaseModel): class ChatInput(BaseModel): - """Input model containing the user's message and conversation history.""" + """Input model containing the user's message.""" - conversation_history: Optional[list[Message]] = Field( - default=None, - description="Previous messages in the conversation, if any", - ) user_message: str = Field( description="The current message from the user", ) @@ -104,7 +80,7 @@ class ChatInput(BaseModel): id="ecommerce-chatbot", model=Model.LLAMA_3_3_70B, ) -async def get_product_recommendations(chat_input: ChatInput) -> Run[ChatbotOutput]: +async def get_product_recommendations(chat_input: ChatInput) -> ChatbotOutput: """ Act as a knowledgeable e-commerce shopping assistant. @@ -142,63 +118,34 @@ async def main(): print("\nExample 1: Looking for headphones") print("-" * 50) - chat_input = ChatInput( - user_message="I'm looking for noise-cancelling headphones for travel. My budget is around $300.", + run = await get_product_recommendations.run( + ChatInput(user_message="I'm looking for noise-cancelling headphones for travel. My budget is around $300."), ) - - run = await get_product_recommendations(chat_input) print(run) - # Example 2: Follow-up question with conversation history + # Example 2: Follow-up question using reply print("\nExample 2: Follow-up about battery life") print("-" * 50) - chat_input = ChatInput( - user_message="Which one has the best battery life?", - conversation_history=[ - Message( - role=Role.USER, - content="I'm looking for noise-cancelling headphones for travel. My budget is around $300.", - ), - run.output.assistant_message, - ], - ) - - run = await get_product_recommendations(chat_input) + run = await run.reply(user_message="Which one has the best battery life?") print(run) # Example 3: Specific question about a previously recommended product print("\nExample 3: Question about a specific product") print("-" * 50) - chat_input = ChatInput( - user_message="Tell me more about the noise cancellation features of the first headphone you recommended.", - conversation_history=[ - Message( - role=Role.USER, - content="I'm looking for noise-cancelling headphones for travel. My budget is around $300.", - ), - run.output.assistant_message, - Message( - role=Role.USER, - content="Which one has the best battery life?", - ), - run.output.assistant_message, - ], + run = await run.reply( + user_message=( + "Tell me more about the noise cancellation features of the first headphone you recommended." + ), ) - - run = await get_product_recommendations(chat_input) print(run) # Example 4: Different product category print("\nExample 4: Looking for a TV") print("-" * 50) - chat_input = ChatInput( - user_message="I need a good TV for gaming. My budget is $1000.", - ) - - run = await get_product_recommendations(chat_input) + run = await run.reply(user_message="I need a good TV for gaming. My budget is $1000.") print(run) diff --git a/examples/13_rag.py b/examples/13_rag.py index ae2f851..08ac40b 100644 --- a/examples/13_rag.py +++ b/examples/13_rag.py @@ -2,7 +2,7 @@ This example demonstrates how to create a RAG-enabled chatbot that: 1. Uses a search tool to find relevant information from a knowledge base 2. Incorporates search results into its responses -3. Maintains conversation context +3. Maintains context through conversation using .reply 4. Provides well-structured, informative responses Note: WorkflowAI does not manage the RAG implementation (yet). You need to provide your own @@ -11,20 +11,11 @@ """ import asyncio -from enum import Enum -from typing import Optional from pydantic import BaseModel, Field import workflowai -from workflowai import Model, Run - - -class Role(str, Enum): - """Enum representing possible message roles.""" - - USER = "user" - ASSISTANT = "assistant" +from workflowai import Model class SearchResult(BaseModel): @@ -79,38 +70,25 @@ async def search_faq(query: str) -> list[SearchResult]: ] -class Message(BaseModel): - """Model representing a chat message.""" +class AssistantMessage(BaseModel): + """Model representing a message from the assistant.""" - role: Role content: str = Field( description="The content of the message", ) -class AssistantMessage(Message): - """Model representing a message from the assistant.""" - - role: Role = Role.ASSISTANT - content: str = "" - - class ChatbotOutput(BaseModel): """Output model for the chatbot response.""" assistant_message: AssistantMessage = Field( - default_factory=AssistantMessage, description="The chatbot's response message", ) class ChatInput(BaseModel): - """Input model containing the user's message and conversation history.""" + """Input model containing the user's message.""" - conversation_history: Optional[list[Message]] = Field( - default=None, - description="Previous messages in the conversation, if any", - ) user_message: str = Field( description="The current message from the user", ) @@ -124,7 +102,7 @@ class ChatInput(BaseModel): # The agent will automatically handle calling the tool and incorporating the results. tools=[search_faq], ) -async def chat_agent(chat_input: ChatInput) -> Run[ChatbotOutput]: +async def chat_agent(chat_input: ChatInput) -> ChatbotOutput: """ Act as a knowledgeable assistant that uses search to find and incorporate relevant information. @@ -165,15 +143,27 @@ async def chat_agent(chat_input: ChatInput) -> Run[ChatbotOutput]: async def main(): - # Example: Question about return policy - print("\nExample: Question about return policy") + # Example 1: Initial question about return policy + print("\nExample 1: Question about return policy") print("-" * 50) - chat_input = ChatInput( - user_message="What is your return policy? Can I return items I bought online?", + run = await chat_agent.run( + ChatInput(user_message="What is your return policy? Can I return items I bought online?"), ) + print(run) + + # Example 2: Follow-up question about shipping + print("\nExample 2: Follow-up about shipping") + print("-" * 50) + + run = await run.reply(user_message="How long does shipping usually take?") + print(run) + + # Example 3: Specific question about refund processing + print("\nExample 3: Question about refunds") + print("-" * 50) - run = await chat_agent(chat_input) + run = await run.reply(user_message="Once I return an item, how long until I get my refund?") print(run) diff --git a/examples/15_pii_extraction.py b/examples/15_pii_extraction.py new file mode 100644 index 0000000..c1f069f --- /dev/null +++ b/examples/15_pii_extraction.py @@ -0,0 +1,137 @@ +""" +This example demonstrates how to create an agent that extracts and redacts Personal Identifiable +Information (PII) from text. It showcases: + +1. Handling sensitive information with clear categorization +2. Structured output with both redacted text and extracted PII +3. Enum usage for PII categories +4. Comprehensive PII detection and redaction +""" + +import asyncio +from enum import Enum + +from pydantic import BaseModel, Field + +import workflowai +from workflowai import Model + + +class PIIType(str, Enum): + """Categories of Personal Identifiable Information.""" + NAME = "NAME" # Full names, first names, last names + EMAIL = "EMAIL" # Email addresses + PHONE = "PHONE" # Phone numbers, fax numbers + ADDRESS = "ADDRESS" # Physical addresses, postal codes + SSN = "SSN" # Social Security Numbers, National IDs + DOB = "DOB" # Date of birth, age + FINANCIAL = "FINANCIAL" # Credit card numbers, bank accounts + LICENSE = "LICENSE" # Driver's license, professional licenses + URL = "URL" # Personal URLs, social media profiles + OTHER = "OTHER" # Other types of PII not covered above + + +class PIIExtraction(BaseModel): + """Represents an extracted piece of PII with its type.""" + text: str = Field(description="The extracted PII text") + type: PIIType = Field(description="The category of PII") + start_index: int = Field(description="Starting position in the original text") + end_index: int = Field(description="Ending position in the original text") + + +class PIIInput(BaseModel): + """Input model for PII extraction.""" + text: str = Field( + description="The text to analyze for PII", + examples=[ + "Hi, I'm John Doe. You can reach me at john.doe@email.com or call 555-0123. " + "My SSN is 123-45-6789 and I live at 123 Main St, Springfield, IL 62701.", + ], + ) + + +class PIIOutput(BaseModel): + """Output model containing redacted text and extracted PII.""" + redacted_text: str = Field( + description="The original text with all PII replaced by [REDACTED]", + examples=[ + "Hi, I'm [REDACTED]. You can reach me at [REDACTED] or call [REDACTED]. " + "My SSN is [REDACTED] and I live at [REDACTED].", + ], + ) + extracted_pii: list[PIIExtraction] = Field( + description="List of extracted PII items with their types and positions", + examples=[ + [ + {"text": "John Doe", "type": "NAME", "start_index": 8, "end_index": 16}, + {"text": "john.doe@email.com", "type": "EMAIL", "start_index": 30, "end_index": 47}, + {"text": "555-0123", "type": "PHONE", "start_index": 57, "end_index": 65}, + ], + ], + ) + + +@workflowai.agent( + id="pii-extractor", + model=Model.CLAUDE_3_5_SONNET_LATEST, +) +async def extract_pii(input_data: PIIInput) -> PIIOutput: + """ + Extract and redact Personal Identifiable Information (PII) from text. + + Guidelines: + 1. Identify all instances of PII in the input text + 2. Categorize each PII instance into one of the defined types + 3. Record the exact position (start and end indices) of each PII instance + 4. Replace all PII in the text with [REDACTED] + 5. Ensure no sensitive information is left unredacted + 6. Be thorough but avoid over-redacting non-PII information + 7. When in doubt about PII type, use the OTHER category + 8. Maintain the original text structure and formatting + 9. Handle overlapping PII appropriately (e.g., name within an email) + 10. Consider context when identifying PII (e.g., distinguish between company and personal emails) + """ + ... + + +async def main(): + # Example 1: Basic PII extraction + print("\nExample 1: Basic PII") + print("-" * 50) + text = ( + "Hello, my name is Sarah Johnson and my email is sarah.j@example.com. " + "You can reach me at (555) 123-4567 or visit my blog at blog.sarahj.net. " + "I was born on 03/15/1985." + ) + result = await extract_pii.run(PIIInput(text=text)) + print("\nOriginal text:") + print(text) + print("\nRedacted text:") + print(result.output.redacted_text) + print("\nExtracted PII:") + for pii in result.output.extracted_pii: + print(f"- {pii.type}: {pii.text} (positions {pii.start_index}-{pii.end_index})") + + # Example 2: Complex PII with financial and address information + print("\n\nExample 2: Complex PII") + print("-" * 50) + text = ( + "Customer: David Wilson\n" + "Card: 4532-9678-1234-5678\n" + "Address: 789 Oak Avenue, Apt 4B\n" + " Boston, MA 02108\n" + "License: MA12-345-678\n" + "SSN: 078-05-1120" + ) + result = await extract_pii.run(PIIInput(text=text)) + print("\nOriginal text:") + print(text) + print("\nRedacted text:") + print(result.output.redacted_text) + print("\nExtracted PII:") + for pii in result.output.extracted_pii: + print(f"- {pii.type}: {pii.text} (positions {pii.start_index}-{pii.end_index})") + + +if __name__ == "__main__": + asyncio.run(main())