import requests
from typing import Dict, List, Callable
from observability import span, current_observation
from langchain.prompts.example_selector.base import BaseExampleSelector


class EmbeddingExampleSelector(BaseExampleSelector):
    def __init__(self, embedding_endpoint: str, namespace: str, input_fn: Callable = None, output_fn: Callable = None):
        super(EmbeddingExampleSelector, self).__init__()
        self.embedding_endpoint = embedding_endpoint
        self.namespace = namespace
        self.input_fn = input_fn
        self.output_fn = output_fn

    def add_example(self, example: Dict[str, str]) -> None:
        """Add new example to store for a key."""
        
        with requests.put(f'{self.embedding_endpoint}/{self.namespace}/upsert', json=example) as r:
            print(r.text)
            if r.status_code != 200:
                raise Exception(f'add example error: {r.text}')
    
    @span(name="select examples")
    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
        current_observation().update(metadata={
            "namespace": self.namespace,
        })

        if self.input_fn is not None:
            input_variables = self.input_fn(input_variables)

        with requests.post(
            f'{self.embedding_endpoint}/{self.namespace}/query',
            json=input_variables,
        ) as r:
            x = r.json()
            if x.get('status') != 'success':
                raise Exception('embedding service call not success')
                
            if self.output_fn is not None:
                return list(map(self.output_fn, x['matches']))
            return x['matches']