#!/usr/bin/env python
"""A more complex example that shows how to configure index name at run time."""
from typing import Any, Iterable, List, Optional, Type

from fastapi import FastAPI
from langchain.embeddings import OpenAIEmbeddings
from langchain.schema import Document
from langchain.schema.embeddings import Embeddings
from langchain.schema.retriever import BaseRetriever
from langchain.schema.runnable import (
    ConfigurableFieldSingleOption,
    RunnableConfig,
    RunnableSerializable,
)
from langchain.schema.vectorstore import VST
from langchain.vectorstores import FAISS, VectorStore

from langserve import add_routes
from langserve.pydantic_v1 import BaseModel, Field

vectorstore1 = FAISS.from_texts(
    ["cats like fish", "dogs like sticks"], embedding=OpenAIEmbeddings()
)

vectorstore2 = FAISS.from_texts(["x_n+1=a * xn * (1-xn)"], embedding=OpenAIEmbeddings())


app = FastAPI(
    title="LangChain Server",
    version="1.0",
    description="Spin up a simple api server using Langchain's Runnable interfaces",
)


class UnderlyingVectorStore(VectorStore):
    """This is a fake vectorstore for demo purposes."""

    def __init__(self, collection_name: str) -> None:
        """Fake vectorstore that has a collection name."""
        self.collection_name = collection_name

    def as_retriever(self) -> BaseRetriever:
        if self.collection_name == "index1":
            return vectorstore1.as_retriever()
        elif self.collection_name == "index2":
            return vectorstore2.as_retriever()
        else:
            raise NotImplementedError(
                f"No retriever for collection {self.collection_name}"
            )

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ) -> List[str]:
        raise NotImplementedError()

    @classmethod
    def from_texts(
        cls: Type[VST],
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ) -> VST:
        raise NotImplementedError()

    def similarity_search(
        self, embedding: List[float], k: int = 4, **kwargs: Any
    ) -> List[Document]:
        raise NotImplementedError()


class ConfigurableRetriever(RunnableSerializable[str, List[Document]]):
    """Create a custom retriever that can be configured by the user.

    This is an example of how to create a custom runnable that can be configured
    to use a different collection name at run time.

    Configuration involves instantiating a VectorStore with a collection name.
    at run time, so the underlying vectorstore should be *cheap* to instantiate.

    For example, it should not be making any network requests at instantiation time.

    Make sure that the vectorstore you use meets this criteria.
    """

    collection_name: str

    def invoke(
        self, input: str, config: Optional[RunnableConfig] = None
    ) -> List[Document]:
        """Invoke the retriever."""
        vectorstore = UnderlyingVectorStore(self.collection_name)
        retriever = vectorstore.as_retriever()
        return retriever.invoke(input, config=config)


configurable_collection_name = ConfigurableRetriever(
    collection_name="index1"
).configurable_fields(
    collection_name=ConfigurableFieldSingleOption(
        id="collection_name",
        name="Collection Name",
        description="The name of the collection to use for the retriever.",
        options={
            "Index 1": "index1",
            "Index 2": "index2",
        },
        default="Index 1",
    )
)


class Request(BaseModel):
    __root__: str = Field(default="cat", description="Search query")


add_routes(app, configurable_collection_name.with_types(input_type=Request))

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="localhost", port=8000)

langchain.schema.runnable.ConfigurableFieldSingleOption langchain.embeddings.OpenAIEmbeddings
[((893, 1027), 'fastapi.FastAPI', 'FastAPI', ([], {'title': '"""LangChain Server"""', 'version': '"""1.0"""', 'description': '"""Spin up a simple api server using Langchain\'s Runnable interfaces"""'}), '(title=\'LangChain Server\', version=\'1.0\', description=\n    "Spin up a simple api server using Langchain\'s Runnable interfaces")\n', (900, 1027), False, 'from fastapi import FastAPI\n'), ((3661, 3709), 'langserve.pydantic_v1.Field', 'Field', ([], {'default': '"""cat"""', 'description': '"""Search query"""'}), "(default='cat', description='Search query')\n", (3666, 3709), False, 'from langserve.pydantic_v1 import BaseModel, Field\n'), ((3841, 3886), 'uvicorn.run', 'uvicorn.run', (['app'], {'host': '"""localhost"""', 'port': '(8000)'}), "(app, host='localhost', port=8000)\n", (3852, 3886), False, 'import uvicorn\n'), ((774, 792), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (790, 792), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((865, 883), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (881, 883), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((3314, 3536), 'langchain.schema.runnable.ConfigurableFieldSingleOption', 'ConfigurableFieldSingleOption', ([], {'id': '"""collection_name"""', 'name': '"""Collection Name"""', 'description': '"""The name of the collection to use for the retriever."""', 'options': "{'Index 1': 'index1', 'Index 2': 'index2'}", 'default': '"""Index 1"""'}), "(id='collection_name', name='Collection Name',\n    description='The name of the collection to use for the retriever.',\n    options={'Index 1': 'index1', 'Index 2': 'index2'}, default='Index 1')\n", (3343, 3536), False, 'from langchain.schema.runnable import ConfigurableFieldSingleOption, RunnableConfig, RunnableSerializable\n')]