Infrastructure Setup

import os

ENV = os.getenv('ENVIRONMENT', 'dev')
APP_NAMESPACE = f'{ENV}_myapp'

# Model Configuration
EMBEDDING_MODEL = os.getenv('EMBEDDING_MODEL', 'intfloat/e5-large-v2')
OPENAI_MODEL = os.getenv('OPENAI_MODEL', 'gpt-4o-mini')

# Storage
MEDIA_STORAGE_BUCKET = os.getenv('MEDIA_STORAGE_BUCKET')

# Prompts
RAG_SYSTEM_PROMPT = """You are a helpful assistant. Use the provided context to answer questions."""

import pixeltable as pxt

@pxt.udf
def format_prompt(context: list, question: str) -> str:
    """Format RAG prompt with context."""
    context_str = "\n".join([doc['text'] for doc in context])
    return f"Context:\n{context_str}\n\nQuestion: {question}"

@pxt.udf(resource_pool='request-rate:my_service')
async def call_custom_model(prompt: str) -> dict:
    """Call self-hosted model endpoint."""
    # Your custom logic here
    return {"response": "..."}

import pixeltable as pxt
from pixeltable.functions.huggingface import sentence_transformer
import config
from functions import format_prompt  # Import module UDFs

# Create namespace
pxt.create_dir(config.APP_NAMESPACE, if_exists='ignore')

# Define base table
docs = pxt.create_table(
    f'{config.APP_NAMESPACE}/documents',
    {'document': pxt.Document, 'metadata': pxt.Json, 'timestamp': pxt.Timestamp},
    if_exists='ignore'
)

# Add computed columns
docs.add_computed_column(
    embedding=sentence_transformer(docs.document, model_id=config.EMBEDDING_MODEL),
    if_exists='ignore'
)

# Add embedding index for similarity search
docs.add_embedding_index('embedding', metric='cosine', if_not_exists=True)

# Define retrieval query function
@pxt.query
def search_documents(query_text: str, limit: int = 5):
    """RAG retrieval query."""
    sim = docs.embedding.similarity(string=query_text)
    return docs.order_by(sim, asc=False).limit(limit).select(docs.document, sim)

from fastapi import FastAPI
import pixeltable as pxt
from setup_pixeltable import search_documents
import config

app = FastAPI()
docs_table = pxt.get_table(f'{config.APP_NAMESPACE}/documents')

@app.get("/search")
async def search(query: str, limit: int = 5):
    results = search_documents(query, limit).collect()
    return {"results": list(results)}