Working with Bedrock in Pixeltable

%pip install -qU pixeltable boto3 duckduckgo-search

import os
import getpass

if 'AWS_ACCESS_KEY' not in os.environ:
    os.environ['AWS_ACCESS_KEY'] = getpass.getpass('Enter your AWS Access Key:')

if 'AWS_SECRET_ACCESS_KEY' not in os.environ:
    os.environ['AWS_SECRET_ACCESS_KEY'] = getpass.getpass('Enter your AWS Secret Access Key:')

import pixeltable as pxt

# Remove the `bedrock_demo` directory and its contents, if it exists
pxt.drop_dir('bedrock_demo', force=True)
pxt.create_dir('bedrock_demo')

Created directory 'bedrock_demo'.

<pixeltable.catalog.dir.Dir at 0x1d4253ae930>

from pixeltable.functions import bedrock

# Create a table in Pixeltable and pick a model hosted on Bedrock with some parameters

t = pxt.create_table('bedrock_demo.chat', {'input': pxt.String})

msgs = [{'role': 'user', 'content': t.input}]
t.add_computed_column(output=bedrock.converse(
        model_id="amazon.nova-pro-v1:0",
        messages=[
            {
                'role': 'user',
                'content': [
                    {
                        'text': t.input,
                    }
                ]
            }
        ],
    ))

Created table `chat`.
Added 0 column values with 0 errors.

UpdateStatus(num_rows=0, num_computed_values=0, num_excs=0, updated_cols=[], cols_with_excs=[])

# Parse the response into a new column
t.add_computed_column(response=t.output.output.message.content[0].text)

Added 0 column values with 0 errors.

UpdateStatus(num_rows=0, num_computed_values=0, num_excs=0, updated_cols=[], cols_with_excs=[])

# Start a conversation
t.insert([{'input': 'What was the outcome of the 1904 US Presidential election?'}])
t.select(t.input, t.response).show()

Inserting rows into `chat`: 1 rows [00:00, 147.54 rows/s]
Inserted 1 row with 0 errors.

import pixeltable as pxt
import pixeltable.functions as pxtf
from pixeltable.functions.bedrock import converse, invoke_tools
from duckduckgo_search import DDGS

# Initialize app structure
pxt.drop_dir("agents", force=True)
pxt.create_dir("agents")

Created directory 'agents'.

<pixeltable.catalog.dir.Dir at 0x1d4252acdd0>

# Define tools
@pxt.udf
def search_news(keywords: str, max_results: int) -> str:
    """Search news using DuckDuckGo and return results."""
    try:
        with DDGS() as ddgs:
            results = ddgs.news(
                keywords=keywords,
                region="wt-wt",
                safesearch="off",
                timelimit="m",
                max_results=max_results,
            )
            formatted_results = []
            for i, r in enumerate(results, 1):
                formatted_results.append(
                    f"{i}. Title: {r['title']}\n"
                    f"   Source: {r['source']}\n"
                    f"   Published: {r['date']}\n"
                    f"   Snippet: {r['body']}\n"
                )
            return "\n".join(formatted_results)
    except Exception as e:
        return f"Search failed: {str(e)}"

@pxt.udf
def get_weather(location: str) -> str:
    """Mock weather function - replace with actual API call."""
    return f"Current weather in {location}: 72°F, Partly Cloudy"

# Register all tools
tools = pxt.tools(search_news, get_weather)

# Create base table
tool_agent = pxt.create_table(
    "agents.tools", 
    {"prompt": pxt.String}, 
    if_exists="ignore"
)

# Add tool selection and execution workflow
tool_agent.add_computed_column(
    initial_response=converse(
        model_id="amazon.nova-pro-v1:0",
        messages=[
            {
                'role': 'user',
                'content': [
                    {
                        'text': tool_agent.prompt,
                    }
                ]
            }
        ],
        tool_config=tools,
    )
)

# Add tool execution
tool_agent.add_computed_column(
    tool_output=invoke_tools(tools, tool_agent.initial_response)
)

Created table `tools`.
Added 0 column values with 0 errors.
Added 0 column values with 0 errors.

UpdateStatus(num_rows=0, num_computed_values=0, num_excs=0, updated_cols=[], cols_with_excs=[])

# Add response formatting
tool_agent.add_computed_column(
    tool_response_prompt=pxtf.string.format(
        "Orginal Prompt\n{0}: Tool Output\n{1}", 
        tool_agent.prompt, 
        tool_agent.tool_output
    ),
    if_exists="ignore",
)

# Add final response generation
tool_agent.add_computed_column(
    final_response=converse(
        model_id="amazon.nova-pro-v1:0",
        messages=[
            {
                'role': 'user',
                'content': [
                    {
                        'text': tool_agent.tool_response_prompt,
                    }
                ]
            }
        ]
    )
)

tool_agent.add_computed_column(
    answer=tool_agent.final_response.output.message.content[0].text
)

Added 0 column values with 0 errors.
Added 0 column values with 0 errors.
Added 0 column values with 0 errors.

UpdateStatus(num_rows=0, num_computed_values=0, num_excs=0, updated_cols=[], cols_with_excs=[])

# Example queries using different tools
queries = [
    "What's the latest news about SpaceX?",
    "What's the weather in San Francisco?",
]

# Use the agent
for query in queries:
    tool_agent.insert([{'prompt': query}])
    result = tool_agent.select(
        tool_agent.prompt,
        tool_agent.tool_output,
        tool_agent.answer
    ).tail(1)
    print(f"\nQuery: {query}")
    print(f"Answer: {result['answer'][0]}")

# Display the full table
tool_agent.select(tool_agent.prompt, tool_agent.answer).show()

Inserting rows into `tools`: 1 rows [00:00, 168.89 rows/s]
Inserted 1 row with 0 errors.

Query: What's the latest news about SpaceX?
Answer: Here's the latest news about SpaceX:

1. **Starbase City Proposal**:
   - **Title**: The home of Elon Musk's SpaceX could become an official Texas city called Starbase
   - **Source**: The Associated Press
   - **Published**: 2025-04-30
   - **Snippet**: An election is scheduled for Saturday to determine whether a small area of coastal South Texas, which is home to SpaceX, will become an official city named Starbase.

2. **250th Starlink Satellite Mission**:
   - **Title**: SpaceX launches 250th Starlink satellite mission, lands rocket at sea (video, photos)
   - **Source**: Space.com
   - **Published**: 2025-04-28
   - **Snippet**: SpaceX achieved a milestone by launching its 250th Starlink mission. A Falcon 9 rocket carrying 23 Starlink satellites, including 13 with direct-to-cell capability, lifted off from Cape Canaveral Space Force Station.

3. **Back-to-Back Starlink Launches**:
   - **Title**: SpaceX Falcon 9 rocket launches 1st of 2 planned Starlink launches in 2 days, lands booster at sea (video)
   - **Source**: Space.com
   - **Published**: 2025-04-28
   - **Snippet**: SpaceX conducted its 49th Falcon 9 mission of 2025 and has another Starlink launch scheduled for later that night from Cape Canaveral Space Force Station.

4. **Amazon’s Satellite Internet Competition**:
   - **Title**: Amazon Just Launched Its First Internet Satellites to Compete Against SpaceX's Starlink
   - **Source**: The Associated Press
   - **Published**: 2025-04-29
   - **Snippet**: Amazon launched its first set of internet satellites using an Atlas V rocket, aiming to compete with SpaceX's Starlink.

5. **Recent Starlink Satellite Launch**:
   - **Title**: SpaceX launches batch of Starlink satellites into low-Earth orbit
   - **Source**: UPI
   - **Published**: 2025-04-28
   - **Snippet**: SpaceX launched a Falcon 9 rocket carrying nearly two dozen Starlink satellites into low-Earth orbit from Florida's east coast.
Inserting rows into `tools`: 1 rows [00:00, 363.87 rows/s]
Inserted 1 row with 0 errors.

Query: What's the weather in San Francisco?
Answer: Based on the tool output provided, here's the current weather in San Francisco:

- **Temperature:** 72°F
- **Condition:** Partly Cloudy

If you need more detailed information or a forecast for the upcoming days, feel free to ask!