AI Engineer Skill

"""
Anthropic Claude API Integration with retry logic and monitoring
Production-ready wrapper for Anthropic API calls
"""

import os
import time
import json
import logging
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass
from pathlib import Path
import yaml

try:
    import anthropic
    from anthropic import Anthropic, APIError, RateLimitError, APITimeoutError
except ImportError:
    raise ImportError("Anthropic package required: pip install anthropic")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class AnthropicConfig:
    api_key: str
    model: str = "claude-3-5-sonnet-20241022"
    max_retries: int = 3
    retry_delay: float = 1.0
    timeout: int = 120
    max_tokens: int = 4096
    temperature: float = 0.7
    rate_limit_delay: float = 0.5

    @classmethod
    def from_yaml(cls, path: Union[str, Path]) -> 'AnthropicConfig':
        with open(path, 'r') as f:
            config = yaml.safe_load(f)
        return cls(**config)


class AnthropicIntegration:
    def __init__(self, config: AnthropicConfig):
        self.config = config
        self.client = Anthropic(
            api_key=config.api_key,
            max_retries=config.max_retries,
            timeout=config.timeout
        )
        self.usage_stats = {
            'total_tokens': 0,
            'input_tokens': 0,
            'output_tokens': 0,
            'total_requests': 0,
            'successful_requests': 0,
            'failed_requests': 0
        }

    def messages(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> Dict[str, Any]:
        for attempt in range(self.config.max_retries):
            try:
                time.sleep(self.config.rate_limit_delay)

                response = self.client.messages.create(
                    model=kwargs.get('model', self.config.model),
                    max_tokens=kwargs.get('max_tokens', self.config.max_tokens),
                    temperature=kwargs.get('temperature', self.config.temperature),
                    messages=messages,
                    **{k: v for k, v in kwargs.items()
                       if k not in ['model', 'max_tokens', 'temperature']}
                )

                content = response.content[0].text
                self._update_usage(response.usage)
                self.usage_stats['successful_requests'] += 1
                self.usage_stats['total_requests'] += 1

                return {
                    'content': content,
                    'usage': {
                        'input_tokens': response.usage.input_tokens,
                        'output_tokens': response.usage.output_tokens,
                        'total_tokens': response.usage.input_tokens + response.usage.output_tokens
                    },
                    'model': response.model,
                    'stop_reason': response.stop_reason
                }

            except RateLimitError as e:
                logger.warning(f"Rate limit hit (attempt {attempt + 1}): {e}")
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

            except APITimeoutError as e:
                logger.warning(f"Timeout (attempt {attempt + 1}): {e}")
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

            except APIError as e:
                logger.error(f"API error: {e}")
                self.usage_stats['failed_requests'] += 1
                self.usage_stats['total_requests'] += 1
                raise

    def streaming_messages(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ):
        for attempt in range(self.config.max_retries):
            try:
                time.sleep(self.config.rate_limit_delay)

                with self.client.messages.stream(
                    model=kwargs.get('model', self.config.model),
                    max_tokens=kwargs.get('max_tokens', self.config.max_tokens),
                    temperature=kwargs.get('temperature', self.config.temperature),
                    messages=messages,
                    **{k: v for k, v in kwargs.items()
                       if k not in ['model', 'max_tokens', 'temperature']}
                ) as stream:
                    for text in stream.text_stream:
                        yield text

                return

            except (RateLimitError, APITimeoutError) as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

    def _update_usage(self, usage: Any):
        if usage:
            self.usage_stats['input_tokens'] += usage.input_tokens
            self.usage_stats['output_tokens'] += usage.output_tokens
            self.usage_stats['total_tokens'] += usage.input_tokens + usage.output_tokens

    def get_usage_stats(self) -> Dict[str, Any]:
        return self.usage_stats.copy()

    def estimate_cost(self, pricing: Dict[str, Dict[str, float]]) -> Dict[str, float]:
        input_cost = (self.usage_stats['input_tokens'] / 1000000) * \
                    pricing.get(self.config.model, {}).get('input', 3.0)
        output_cost = (self.usage_stats['output_tokens'] / 1000000) * \
                     pricing.get(self.config.model, {}).get('output', 15.0)

        return {
            'input_cost': input_cost,
            'output_cost': output_cost,
            'total_cost': input_cost + output_cost
        }


def main():
    config = AnthropicConfig(
        api_key=os.getenv("ANTHROPIC_API_KEY", ""),
        model="claude-3-5-sonnet-20241022"
    )

    integration = AnthropicIntegration(config)

    messages = [
        {"role": "user", "content": "Explain machine learning in one sentence."}
    ]

    response = integration.messages(messages)
    print("Response:", response['content'])
    print("Usage:", response['usage'])
    print("Stats:", integration.get_usage_stats())


if __name__ == "__main__":
    main()

"""
OpenAI API Integration with retry logic, rate limiting, and monitoring
Production-ready wrapper for OpenAI API calls
"""

import os
import time
import json
import logging
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass, asdict
from pathlib import Path
import yaml

try:
    import openai
    from openai import OpenAI, APIError, RateLimitError, APITimeoutError
except ImportError:
    raise ImportError("OpenAI package required: pip install openai")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class OpenAIConfig:
    api_key: str
    model: str = "gpt-4"
    max_retries: int = 3
    retry_delay: float = 1.0
    timeout: int = 120
    organization: Optional[str] = None
    base_url: Optional[str] = None
    max_tokens: int = 4096
    temperature: float = 0.7
    rate_limit_delay: float = 0.5

    @classmethod
    def from_yaml(cls, path: Union[str, Path]) -> 'OpenAIConfig':
        with open(path, 'r') as f:
            config = yaml.safe_load(f)
        return cls(**config)


class OpenAIIntegration:
    def __init__(self, config: OpenAIConfig):
        self.config = config
        self.client = OpenAI(
            api_key=config.api_key,
            organization=config.organization,
            base_url=config.base_url,
            timeout=config.timeout,
            max_retries=config.max_retries
        )
        self.usage_stats = {
            'total_tokens': 0,
            'prompt_tokens': 0,
            'completion_tokens': 0,
            'total_requests': 0,
            'successful_requests': 0,
            'failed_requests': 0
        }

    def chat_completion(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> Dict[str, Any]:
        for attempt in range(self.config.max_retries):
            try:
                time.sleep(self.config.rate_limit_delay)
                
                response = self.client.chat.completions.create(
                    model=kwargs.get('model', self.config.model),
                    messages=messages,
                    max_tokens=kwargs.get('max_tokens', self.config.max_tokens),
                    temperature=kwargs.get('temperature', self.config.temperature),
                    **{k: v for k, v in kwargs.items() 
                       if k not in ['model', 'max_tokens', 'temperature']}
                )

                self._update_usage(response.usage)
                self.usage_stats['successful_requests'] += 1
                self.usage_stats['total_requests'] += 1

                return {
                    'content': response.choices[0].message.content,
                    'usage': {
                        'prompt_tokens': response.usage.prompt_tokens,
                        'completion_tokens': response.usage.completion_tokens,
                        'total_tokens': response.usage.total_tokens
                    },
                    'model': response.model,
                    'finish_reason': response.choices[0].finish_reason
                }

            except RateLimitError as e:
                logger.warning(f"Rate limit hit (attempt {attempt + 1}): {e}")
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

            except APITimeoutError as e:
                logger.warning(f"Timeout (attempt {attempt + 1}): {e}")
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

            except APIError as e:
                logger.error(f"API error: {e}")
                self.usage_stats['failed_requests'] += 1
                self.usage_stats['total_requests'] += 1
                raise

    def streaming_chat_completion(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ):
        for attempt in range(self.config.max_retries):
            try:
                time.sleep(self.config.rate_limit_delay)
                
                stream = self.client.chat.completions.create(
                    model=kwargs.get('model', self.config.model),
                    messages=messages,
                    max_tokens=kwargs.get('max_tokens', self.config.max_tokens),
                    temperature=kwargs.get('temperature', self.config.temperature),
                    stream=True,
                    **{k: v for k, v in kwargs.items() 
                       if k not in ['model', 'max_tokens', 'temperature', 'stream']}
                )

                for chunk in stream:
                    if chunk.choices[0].delta.content is not None:
                        yield chunk.choices[0].delta.content

                return

            except (RateLimitError, APITimeoutError) as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (2 ** attempt))
                    continue
                raise

    def embed_text(
        self,
        texts: List[str],
        model: str = "text-embedding-3-small",
        **kwargs
    ) -> List[List[float]]:
        try:
            response = self.client.embeddings.create(
                model=model,
                input=texts,
                **kwargs
            )
            return [embedding.embedding for embedding in response.data]
        except APIError as e:
            logger.error(f"Embedding error: {e}")
            raise

    def _update_usage(self, usage: Any):
        if usage:
            self.usage_stats['total_tokens'] += usage.total_tokens
            self.usage_stats['prompt_tokens'] += usage.prompt_tokens
            self.usage_stats['completion_tokens'] += usage.completion_tokens

    def get_usage_stats(self) -> Dict[str, Any]:
        return self.usage_stats.copy()

    def estimate_cost(self, pricing: Dict[str, Dict[str, float]]) -> Dict[str, float]:
        input_cost = (self.usage_stats['prompt_tokens'] / 1000) * \
                    pricing.get(self.config.model, {}).get('input', 0.01)
        output_cost = (self.usage_stats['completion_tokens'] / 1000) * \
                     pricing.get(self.config.model, {}).get('output', 0.03)
        
        return {
            'input_cost': input_cost,
            'output_cost': output_cost,
            'total_cost': input_cost + output_cost
        }


def main():
    config = OpenAIConfig(
        api_key=os.getenv("OPENAI_API_KEY", ""),
        model="gpt-4"
    )

    integration = OpenAIIntegration(config)

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Explain machine learning in one sentence."}
    ]

    response = integration.chat_completion(messages)
    print("Response:", response['content'])
    print("Usage:", response['usage'])
    print("Stats:", integration.get_usage_stats())


if __name__ == "__main__":
    main()

"""
Prompt Template Management System
Manages, versions, and retrieves prompt templates
"""

import os
import json
import yaml
import logging
from typing import Dict, List, Any, Optional, Union
from pathlib import Path
from dataclasses import dataclass, asdict
from datetime import datetime
import hashlib

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class PromptTemplate:
    name: str
    template: str
    version: str
    variables: List[str]
    description: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None

    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.utcnow().isoformat()
        if self.updated_at is None:
            self.updated_at = self.created_at

    def render(self, **kwargs) -> str:
        missing_vars = set(self.variables) - set(kwargs.keys())
        if missing_vars:
            raise ValueError(f"Missing required variables: {missing_vars}")

        try:
            return self.template.format(**kwargs)
        except KeyError as e:
            raise ValueError(f"Template rendering error: {e}")


class PromptManager:
    def __init__(self, templates_dir: Union[str, Path] = "./prompt_templates"):
        self.templates_dir = Path(templates_dir)
        self.templates_dir.mkdir(parents=True, exist_ok=True)
        self.templates: Dict[str, List[PromptTemplate]] = {}

        self._load_templates()

    def _load_templates(self):
        for file_path in self.templates_dir.rglob("*.yaml"):
            try:
                with open(file_path, 'r') as f:
                    data = yaml.safe_load(f)

                for template_data in data.get('templates', []):
                    template = PromptTemplate(**template_data)
                    if template.name not in self.templates:
                        self.templates[template.name] = []
                    self.templates[template.name].append(template)

                logger.info(f"Loaded templates from {file_path}")
            except Exception as e:
                logger.error(f"Error loading {file_path}: {e}")

    def add_template(self, template: PromptTemplate, persist: bool = True):
        if template.name not in self.templates:
            self.templates[template.name] = []

        existing_versions = [t.version for t in self.templates[template.name]]
        if template.version in existing_versions:
            logger.warning(f"Version {template.version} already exists for {template.name}")
            template.version = self._get_next_version(template.name)

        template.updated_at = datetime.utcnow().isoformat()
        self.templates[template.name].append(template)

        if persist:
            self._save_template(template)

        logger.info(f"Added template: {template.name} v{template.version}")

    def _get_next_version(self, name: str) -> str:
        versions = [t.version for t in self.templates.get(name, [])]
        if not versions:
            return "1.0.0"

        try:
            latest = max(versions, key=lambda v: [int(x) for x in v.split('.')])
            major, minor, patch = latest.split('.')
            return f"{major}.{minor}.{int(patch) + 1}"
        except:
            return "1.0.0"

    def _save_template(self, template: PromptTemplate):
        file_path = self.templates_dir / f"{template.name}.yaml"

        data = {'templates': []}
        if file_path.exists():
            with open(file_path, 'r') as f:
                data = yaml.safe_load(f) or {'templates': []}

        template_dict = asdict(template)
        data['templates'].append(template_dict)

        with open(file_path, 'w') as f:
            yaml.dump(data, f, default_flow_style=False)

    def get_template(
        self,
        name: str,
        version: Optional[str] = None
    ) -> Optional[PromptTemplate]:
        if name not in self.templates:
            return None

        if version:
            for template in self.templates[name]:
                if template.version == version:
                    return template
        else:
            return self.templates[name][-1]

        return None

    def list_templates(self, name: Optional[str] = None) -> List[Dict[str, Any]]:
        if name:
            templates = self.templates.get(name, [])
        else:
            templates = []
            for template_list in self.templates.values():
                templates.extend(template_list)

        return [
            {
                'name': t.name,
                'version': t.version,
                'description': t.description,
                'created_at': t.created_at,
                'updated_at': t.updated_at
            }
            for t in templates
        ]

    def delete_template(self, name: str, version: Optional[str] = None):
        if name not in self.templates:
            return

        if version:
            self.templates[name] = [t for t in self.templates[name] if t.version != version]
        else:
            del self.templates[name]

        logger.info(f"Deleted template: {name}")


def create_default_templates(manager: PromptManager):
    templates = [
        PromptTemplate(
            name="code_explanation",
            template="Explain the following code:\n\n```\n{code}\n```\n\nFocus on: {focus}",
            variables=["code", "focus"],
            description="Explains code with specific focus areas",
            metadata={"category": "code"}
        ),
        PromptTemplate(
            name="summarization",
            template="Summarize the following text in {max_sentences} sentences:\n\n{text}",
            variables=["text", "max_sentences"],
            description="Summarizes text to specified length",
            metadata={"category": "nlp"}
        ),
        PromptTemplate(
            name="question_answering",
            template="Based on the following context:\n\n{context}\n\nAnswer the question: {question}",
            variables=["context", "question"],
            description="Answers questions based on provided context",
            metadata={"category": "qa"}
        )
    ]

    for template in templates:
        manager.add_template(template)


def main():
    manager = PromptManager()

    create_default_templates(manager)

    template = manager.get_template("code_explanation")
    if template:
        rendered = template.render(
            code="print('Hello, World!')",
            focus="functionality and best practices"
        )
        print("Rendered prompt:")
        print(rendered)

    print("\nAvailable templates:")
    for t in manager.list_templates():
        print(f"  - {t['name']} v{t['version']}: {t['description']}")


if __name__ == "__main__":
    main()

"""
Token Usage and Cost Optimization
Tracks token usage, estimates costs, and suggests optimizations
"""

import logging
from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass, field
from collections import defaultdict
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class TokenPricing:
    model: str
    input_price_per_1k: float
    output_price_per_1k: float


@dataclass
class TokenUsage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class UsageRecord:
    model: str
    usage: TokenUsage
    timestamp: str
    metadata: Optional[Dict[str, Any]] = None


class TokenTracker:
    def __init__(self):
        self.records: List[UsageRecord] = []
        self.pricing: Dict[str, TokenPricing] = {
            'gpt-4': TokenPricing('gpt-4', 0.03, 0.06),
            'gpt-4-turbo': TokenPricing('gpt-4-turbo', 0.01, 0.03),
            'gpt-3.5-turbo': TokenPricing('gpt-3.5-turbo', 0.0005, 0.0015),
            'claude-3-5-sonnet-20241022': TokenPricing('claude-3-5-sonnet-20241022', 0.003, 0.015),
            'claude-3-opus-20240229': TokenPricing('claude-3-opus-20240229', 0.015, 0.075)
        }

    def record_usage(
        self,
        model: str,
        usage: Dict[str, int],
        metadata: Optional[Dict[str, Any]] = None
    ):
        record = UsageRecord(
            model=model,
            usage=TokenUsage(
                prompt_tokens=usage.get('prompt_tokens', 0),
                completion_tokens=usage.get('completion_tokens', 0),
                total_tokens=usage.get('total_tokens', 0)
            ),
            timestamp=__import__('datetime').datetime.utcnow().isoformat(),
            metadata=metadata or {}
        )
        self.records.append(record)
        logger.debug(f"Recorded usage: {model} - {usage.get('total_tokens', 0)} tokens")

    def get_total_tokens(self, model: Optional[str] = None) -> int:
        total = 0
        for record in self.records:
            if model is None or record.model == model:
                total += record.usage.total_tokens
        return total

    def get_total_cost(self, model: Optional[str] = None) -> float:
        total_cost = 0.0

        for record in self.records:
            if model is not None and record.model != model:
                continue

            pricing = self.pricing.get(record.model)
            if pricing:
                input_cost = (record.usage.prompt_tokens / 1000) * pricing.input_price_per_1k
                output_cost = (record.usage.completion_tokens / 1000) * pricing.output_price_per_1k
                total_cost += input_cost + output_cost

        return total_cost

    def get_cost_by_model(self) -> Dict[str, float]:
        costs = defaultdict(float)
        for record in self.records:
            pricing = self.pricing.get(record.model)
            if pricing:
                input_cost = (record.usage.prompt_tokens / 1000) * pricing.input_price_per_1k
                output_cost = (record.usage.completion_tokens / 1000) * pricing.output_price_per_1k
                costs[record.model] += input_cost + output_cost
        return dict(costs)

    def get_average_tokens_per_request(self, model: Optional[str] = None) -> float:
        records = [r for r in self.records if model is None or r.model == model]
        if not records:
            return 0.0
        return sum(r.usage.total_tokens for r in records) / len(records)

    def get_statistics(self) -> Dict[str, Any]:
        total_tokens = self.get_total_tokens()
        total_cost = self.get_total_cost()
        total_requests = len(self.records)

        avg_tokens = 0.0
        if total_requests > 0:
            avg_tokens = total_tokens / total_requests

        return {
            'total_requests': total_requests,
            'total_tokens': total_tokens,
            'total_cost': total_cost,
            'average_tokens_per_request': avg_tokens,
            'cost_by_model': self.get_cost_by_model(),
            'requests_by_model': self._count_by_model()
        }

    def _count_by_model(self) -> Dict[str, int]:
        counts = defaultdict(int)
        for record in self.records:
            counts[record.model] += 1
        return dict(counts)

    def export_report(self, filepath: str):
        report = {
            'statistics': self.get_statistics(),
            'records': [
                {
                    'model': r.model,
                    'usage': {
                        'prompt_tokens': r.usage.prompt_tokens,
                        'completion_tokens': r.usage.completion_tokens,
                        'total_tokens': r.usage.total_tokens
                    },
                    'timestamp': r.timestamp,
                    'metadata': r.metadata
                }
                for r in self.records
            ]
        }

        with open(filepath, 'w') as f:
            json.dump(report, f, indent=2)

        logger.info(f"Exported report to {filepath}")


class TokenOptimizer:
    @staticmethod
    def optimize_prompt(prompt: str, target_tokens: int) -> tuple[str, Dict[str, Any]]:
        approx_tokens = len(prompt.split()) * 1.3

        suggestions = []

        if approx_tokens > target_tokens:
            suggestions.append("Reduce prompt length")
            suggestions.append("Use more concise language")
            suggestions.append("Remove redundant information")
            suggestions.append("Use system prompt for static context")

        return prompt, {
            'estimated_tokens': approx_tokens,
            'target_tokens': target_tokens,
            'optimizations': suggestions
        }

    @staticmethod
    def suggest_model(
        tokens_needed: int,
        budget: float,
        complexity: str = 'medium'
    ) -> Dict[str, Any]:
        models = [
            {'name': 'gpt-3.5-turbo', 'cost_per_1k': 0.002, 'context': 16384, 'quality': 'low'},
            {'name': 'gpt-4-turbo', 'cost_per_1k': 0.04, 'context': 128000, 'quality': 'high'},
            {'name': 'gpt-4', 'cost_per_1k': 0.09, 'context': 8192, 'quality': 'very-high'},
            {'name': 'claude-3-5-sonnet-20241022', 'cost_per_1k': 0.018, 'context': 200000, 'quality': 'high'},
            {'name': 'claude-3-opus-20240229', 'cost_per_1k': 0.09, 'context': 200000, 'quality': 'very-high'}
        ]

        suitable = []
        for model in models:
            if tokens_needed <= model['context']:
                suitable.append(model)

        suitable.sort(key=lambda m: m['cost_per_1k'])

        quality_ranking = {'low': 1, 'medium': 2, 'high': 3, 'very-high': 4}
        min_quality = quality_ranking.get(complexity, 2)

        for model in suitable:
            if quality_ranking.get(model['quality'], 0) >= min_quality:
                estimated_cost = (tokens_needed / 1000) * model['cost_per_1k']
                if estimated_cost <= budget:
                    return {
                        'recommended': model['name'],
                        'estimated_cost': estimated_cost,
                        'reason': 'Best fit for complexity and budget'
                    }

        return {
            'recommended': suitable[0]['name'] if suitable else 'gpt-3.5-turbo',
            'estimated_cost': (tokens_needed / 1000) * suitable[0]['cost_per_1k'] if suitable else 0,
            'reason': 'Best available option'
        }


def main():
    tracker = TokenTracker()

    sample_usages = [
        {'model': 'gpt-4', 'usage': {'prompt_tokens': 100, 'completion_tokens': 50, 'total_tokens': 150}},
        {'model': 'gpt-3.5-turbo', 'usage': {'prompt_tokens': 200, 'completion_tokens': 100, 'total_tokens': 300}},
        {'model': 'gpt-4', 'usage': {'prompt_tokens': 150, 'completion_tokens': 75, 'total_tokens': 225}},
    ]

    for usage in sample_usages:
        tracker.record_usage(**usage)

    print("=== Token Usage Statistics ===")
    stats = tracker.get_statistics()
    print(json.dumps(stats, indent=2))

    print("\n=== Model Suggestion ===")
    suggestion = TokenOptimizer.suggest_model(tokens_needed=5000, budget=0.50, complexity='high')
    print(json.dumps(suggestion, indent=2))


if __name__ == "__main__":
    main()

"""
RAG (Retrieval-Augmented Generation) Setup Script
Handles vector database, document chunking, and retrieval
"""

import os
import logging
from typing import List, Dict, Any, Optional, Union
from pathlib import Path
from dataclasses import dataclass
import json
import yaml

try:
    import chromadb
    from chromadb.config import Settings
except ImportError:
    raise ImportError("chromadb required: pip install chromadb")

try:
    from sentence_transformers import SentenceTransformer
except ImportError:
    raise ImportError("sentence-transformers required: pip install sentence-transformers")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class RAGConfig:
    collection_name: str = "documents"
    embedding_model: str = "all-MiniLM-L6-v2"
    chunk_size: int = 512
    chunk_overlap: int = 50
    persist_directory: str = "./chroma_db"
    top_k: int = 5

    @classmethod
    def from_yaml(cls, path: Union[str, Path]) -> 'RAGConfig':
        with open(path, 'r') as f:
            config = yaml.safe_load(f)
        return cls(**config)


class DocumentChunker:
    def __init__(self, chunk_size: int = 512, overlap: int = 50):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk_text(self, text: str) -> List[str]:
        chunks = []
        start = 0

        while start < len(text):
            end = start + self.chunk_size
            chunk = text[start:end]

            if start > 0:
                chunk = text[max(0, start - self.overlap):end]

            chunks.append(chunk.strip())
            start = end

        return [c for c in chunks if c]

    def chunk_documents(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        chunked_docs = []

        for doc in documents:
            text = doc.get('text', '')
            chunks = self.chunk_text(text)

            for i, chunk in enumerate(chunks):
                chunked_docs.append({
                    'text': chunk,
                    'metadata': {
                        **doc.get('metadata', {}),
                        'chunk_id': i,
                        'source': doc.get('source', 'unknown')
                    },
                    'id': f"{doc.get('id', 'doc')}_{i}"
                })

        return chunked_docs


class RAGSystem:
    def __init__(self, config: RAGConfig):
        self.config = config
        self.chunker = DocumentChunker(config.chunk_size, config.chunk_overlap)

        Path(config.persist_directory).mkdir(parents=True, exist_ok=True)

        self.client = chromadb.PersistentClient(
            path=config.persist_directory,
            settings=Settings(anonymized_telemetry=False)
        )

        self.collection = self.client.get_or_create_collection(
            name=config.collection_name
        )

        logger.info(f"Loading embedding model: {config.embedding_model}")
        self.embedding_model = SentenceTransformer(config.embedding_model)

    def add_documents(self, documents: List[Dict[str, Any]]):
        chunked_docs = self.chunker.chunk_documents(documents)

        texts = [doc['text'] for doc in chunked_docs]
        embeddings = self.embedding_model.encode(texts).tolist()

        self.collection.add(
            embeddings=embeddings,
            documents=texts,
            metadatas=[doc['metadata'] for doc in chunked_docs],
            ids=[doc['id'] for doc in chunked_docs]
        )

        logger.info(f"Added {len(chunked_docs)} chunks to collection")

    def query(
        self,
        query_text: str,
        n_results: Optional[int] = None,
        where: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        if n_results is None:
            n_results = self.config.top_k

        query_embedding = self.embedding_model.encode([query_text]).tolist()

        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=n_results,
            where=where
        )

        retrieved_docs = []
        for i in range(len(results['ids'][0])):
            retrieved_docs.append({
                'id': results['ids'][0][i],
                'text': results['documents'][0][i],
                'metadata': results['metadatas'][0][i],
                'distance': results['distances'][0][i]
            })

        return retrieved_docs

    def delete_documents(self, ids: List[str]):
        self.collection.delete(ids=ids)
        logger.info(f"Deleted {len(ids)} documents")

    def get_collection_stats(self) -> Dict[str, Any]:
        return {
            'count': self.collection.count(),
            'name': self.config.collection_name,
            'embedding_model': self.config.embedding_model
        }

    def clear_collection(self):
        self.client.delete_collection(name=self.config.collection_name)
        self.collection = self.client.create_collection(name=self.config.collection_name)
        logger.info("Collection cleared")


def load_documents_from_directory(directory: Union[str, Path]) -> List[Dict[str, Any]]:
    documents = []
    directory = Path(directory)

    for file_path in directory.rglob("*.txt"):
        with open(file_path, 'r', encoding='utf-8') as f:
            text = f.read()

        documents.append({
            'id': file_path.stem,
            'text': text,
            'metadata': {
                'source': str(file_path),
                'filename': file_path.name
            }
        })

    return documents


def main():
    config = RAGConfig()
    rag = RAGSystem(config)

    sample_docs = [
        {
            'id': 'doc1',
            'text': 'Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data.',
            'metadata': {'category': 'ML', 'source': 'intro'}
        },
        {
            'id': 'doc2',
            'text': 'Deep learning uses neural networks with multiple layers to model complex patterns in data.',
            'metadata': {'category': 'DL', 'source': 'advanced'}
        }
    ]

    rag.add_documents(sample_docs)

    query = "What is machine learning?"
    results = rag.query(query)

    print(f"Query: {query}")
    for i, result in enumerate(results, 1):
        print(f"\nResult {i}:")
        print(f"Text: {result['text']}")
        print(f"Distance: {result['distance']:.4f}")

    print(f"\nStats: {rag.get_collection_stats()}")


if __name__ == "__main__":
    main()