txtai-bedrock.py

import os

from txtai import Embeddings, RAG

# For demo only. Set via environment.
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""
os.environ["ANN_URL"] = "postgresql+psycopg2://postgres:[email protected]/postgres"
os.environ["CLIENT_URL"] = "postgresql+psycopg2://postgres:[email protected]/postgres"

# Input data
data = [
  "US tops 5 million confirmed virus cases",
  "Canada's last fully intact ice shelf has suddenly collapsed, " +
  "forming a Manhattan-sized iceberg",
  "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
  "The National Park Service warns against sacrificing slower friends " +
  "in a bear attack",
  "Maine man wins $1M from $25 lottery ticket",
  "Make huge profits without work, earn up to $100,000 a day"
]

# Build embeddings index
embeddings = Embeddings(content="client", backend="pgvector", path="bedrock/cohere.embed-english-v3")
embeddings.index(data)

# Create and run pipeline
rag = RAG(embeddings, "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", template="""
  Answer the following question using the provided context.

  Question:
  {question}

  Context:
  {context}
""")

rag("What was won?")