Create 12_contextual_retrieval.py

examples/12_contextual_retrieval.py

@@ -0,0 +1,89 @@
+"""
+This example demonstrates how to implement Contextual Retrieval's context generation
+as described by Anthropic (https://www.anthropic.com/news/contextual-retrieval).
+It shows how to generate concise, contextual descriptions for document chunks.
+"""
+
+import asyncio
+
+from pydantic import BaseModel, Field
+
+import workflowai
+from workflowai import Model, Run
+
+
+class ContextGeneratorInput(BaseModel):
+    """Input for generating context for a document chunk."""
+    doc_content: str = Field(
+        description="The full text content of the document",
+    )
+    chunk_content: str = Field(
+        description="The specific chunk of text to generate context for",
+    )
+
+
+class ContextGeneratorOutput(BaseModel):
+    """Output containing the generated context for a chunk."""
+    context: str = Field(
+        description="Generated contextual information for the chunk",
+        examples=[
+            "This chunk is from section 3.2 discussing revenue growth in Q2 2023",
+            "This appears in the methodology section explaining the experimental setup",
+        ],
+    )
+
+
+@workflowai.agent(
+    id="context-generator",
+    model=Model.CLAUDE_3_5_SONNET_LATEST,
+)
+async def generate_chunk_context(context_input: ContextGeneratorInput) -> Run[ContextGeneratorOutput]:
+    """
+    Here is the chunk we want to situate within the whole document.
+    Please give a short succinct context to situate this chunk within the overall document
+    for the purposes of improving search retrieval of the chunk.
+    """
+    ...
+
+
+async def main():
+    # Example: Generate context for a document chunk
+    print("\nGenerating context for document chunk")
+    print("-" * 50)
+
+    # Example document content
+    doc_content = """
+    ACME Corporation (NASDAQ: ACME)
+    Second Quarter 2023 Financial Results
+
+    Executive Summary
+    ACME Corp. delivered strong performance in Q2 2023, with revenue growth
+    exceeding market expectations. The company's strategic initiatives in
+    AI and cloud services continued to drive expansion.
+
+    Financial Highlights
+    The company's revenue grew by 3% over the previous quarter, reaching
+    $323.4M. This growth was primarily driven by our enterprise segment,
+    which saw a 15% increase in cloud service subscriptions.
+
+    Operational Metrics
+    - Customer base expanded to 15,000 enterprise clients
+    - Cloud platform usage increased by 25%
+    - AI solutions adoption rate reached 40%
+    """
+
+    # Example chunk from the Financial Highlights section
+    chunk_content = "The company's revenue grew by 3% over the previous quarter, reaching $323.4M."
+
+    context_input = ContextGeneratorInput(
+        doc_content=doc_content,
+        chunk_content=chunk_content,
+    )
+
+    run = await generate_chunk_context(context_input)
+    print("\nGenerated Context:")
+    print(run.output.context)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())