AI-Powered Data Pipelines: A Deep Dive

┌─────────────┐     ┌───────────────┐     ┌─────────────────┐
│   Source    │────▶│   Ingestion   │────▶│   Processing    │
│   Systems   │     │   Layer       │     │   Engine        │
└─────────────┘     └───────────────┘     └────────┬────────┘
                                                   │
                    ┌───────────────┐              │
                    │    Vector     │◀─────────────┤
                    │    Database   │              │
                    └───────────────┘              │
                                                   ▼
┌─────────────┐     ┌───────────────┐     ┌─────────────────┐
│   Query     │◀────│   LLM         │◀────│   Enriched      │
│   Layer     │     │   Service     │     │   Data Store    │
└─────────────┘     └───────────────┘     └─────────────────┘

from weaviate import WeaviateClient
from openai import OpenAI

client = WeaviateClient(url="http://localhost:8080")
openai = OpenAI()

def create_embedding(text: str) -> list[float]:
    """Generate embedding using OpenAI's embedding model."""
    response = openai.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

def index_document(doc_id: str, content: str, metadata: dict):
    """Index a document with its embedding."""
    embedding = create_embedding(content)
    
    client.collections.get("documents").data.insert(
        properties={
            "doc_id": doc_id,
            "content": content,
            **metadata
        },
        vector=embedding
    )

def semantic_search(query: str, limit: int = 5) -> list[dict]:
    """Find documents semantically similar to the query."""
    query_embedding = create_embedding(query)
    
    results = client.collections.get("documents").query.near_vector(
        near_vector=query_embedding,
        limit=limit,
        return_metadata=["distance"]
    )
    
    return [
        {
            "content": obj.properties["content"],
            "score": 1 - obj.metadata.distance,
            **obj.properties
        }
        for obj in results.objects
    ]

from pydantic import BaseModel
from openai import OpenAI

class EnrichedRecord(BaseModel):
    original_text: str
    category: str
    sentiment: str
    entities: list[str]
    summary: str
    keywords: list[str]

def enrich_with_llm(text: str) -> EnrichedRecord:
    """Use LLM to extract structured information from text."""
    client = OpenAI()
    
    response = client.chat.completions.create(
        model="gpt-4-turbo-preview",
        messages=[
            {
                "role": "system",
                "content": """Analyze the provided text and extract:
                - Category (news, support, feedback, inquiry, other)
                - Sentiment (positive, negative, neutral)
                - Named entities (people, companies, locations)
                - A brief summary (max 50 words)
                - Key topics/keywords (max 5)
                
                Respond in JSON format."""
            },
            {"role": "user", "content": text}
        ],
        response_format={"type": "json_object"}
    )
    
    data = json.loads(response.choices[0].message.content)
    return EnrichedRecord(original_text=text, **data)

import asyncio
from kafka import KafkaConsumer, KafkaProducer
import json

async def process_stream():
    """Continuously process incoming data."""
    consumer = KafkaConsumer(
        'raw-documents',
        bootstrap_servers=['kafka:9092'],
        value_deserializer=lambda m: json.loads(m.decode('utf-8'))
    )
    
    producer = KafkaProducer(
        bootstrap_servers=['kafka:9092'],
        value_serializer=lambda v: json.dumps(v).encode('utf-8')
    )
    
    for message in consumer:
        doc = message.value
        
        # Enrich with LLM
        enriched = enrich_with_llm(doc['content'])
        
        # Index for vector search
        index_document(
            doc_id=doc['id'],
            content=doc['content'],
            metadata=enriched.dict()
        )
        
        # Publish enriched document
        producer.send('enriched-documents', value=enriched.dict())
        
        print(f"Processed document {doc['id']}")

def answer_with_rag(query: str) -> str:
    """Answer using Retrieval-Augmented Generation."""
    # Retrieve relevant context
    relevant_docs = semantic_search(query, limit=3)
    context = "\n\n".join([doc["content"] for doc in relevant_docs])
    
    response = openai.chat.completions.create(
        model="gpt-4-turbo-preview",
        messages=[
            {
                "role": "system",
                "content": f"""Answer based ONLY on the provided context.
                If the context doesn't contain enough information, say so.
                
                Context:
                {context}"""
            },
            {"role": "user", "content": query}
        ]
    )
    
    return response.choices[0].message.content

from pydantic import BaseModel, validator

class ValidatedOutput(BaseModel):
    category: str
    confidence: float
    
    @validator('category')
    def validate_category(cls, v):
        allowed = {'news', 'support', 'feedback', 'inquiry', 'other'}
        if v not in allowed:
            raise ValueError(f'Category must be one of {allowed}')
        return v
    
    @validator('confidence')
    def validate_confidence(cls, v):
        if not 0 <= v <= 1:
            raise ValueError('Confidence must be between 0 and 1')
        return v