llama-index

1
from typing import Any, Optional
2
from llama_index.core.base.base_query_engine import BaseQueryEngine
3
from llama_index.core.bridge.pydantic import Field, BaseModel
4
from llama_index.networks.schema.contributor import (
5
    ContributorQueryRequest,
6
)
7
from pydantic.v1 import BaseSettings, PrivateAttr
8
from fastapi import FastAPI
9

10

11
class ContributorServiceSettings(BaseSettings):
12
    api_version: str = Field(default="v1", description="API version.")
13
    secret: Optional[str] = Field(
14
        default=None, description="JWT secret."
15
    )  # left for future consideration.
16
    # or if user wants to implement their own
17

18
    class Config:
19
        env_file = ".env", ".env.contributor.service"
20

21

22
class ContributorService(BaseModel):
23
    query_engine: Optional[BaseQueryEngine]
24
    config: ContributorServiceSettings
25
    _fastapi: FastAPI = PrivateAttr()
26

27
    class Config:
28
        arbitrary_types_allowed = True
29

30
    def __init__(self, query_engine, config) -> None:
31
        self._fastapi = FastAPI(
32
            version=config.api_version,
33
        )
34

35
        # routes
36
        self._fastapi.add_api_route(path="/api", endpoint=self.index, methods=["GET"])
37
        self._fastapi.add_api_route(
38
            path="/api/query",
39
            endpoint=self.query,
40
            methods=["POST"],
41
        )
42

43
        super().__init__(query_engine=query_engine, config=config)
44

45
    async def index(self):
46
        """Index endpoint logic."""
47
        return {"message": "Hello World!"}
48

49
    async def query(self, request: ContributorQueryRequest):
50
        """Query endpoint logic."""
51
        result = await self.query_engine.aquery(request.query)
52
        return {
53
            "response": result.response,
54
            "source_nodes": result.source_nodes,
55
            "metadata": result.metadata,
56
        }
57

58
    @classmethod
59
    def from_config_file(
60
        cls, env_file: str, query_engine: BaseQueryEngine
61
    ) -> "ContributorService":
62
        config = ContributorServiceSettings(_env_file=env_file)
63
        return cls(query_engine=query_engine, config=config)
64

65
    def __getattr__(self, attr) -> Any:
66
        if hasattr(self._fastapi, attr):
67
            return getattr(self._fastapi, attr)
68
        else:
69
            raise AttributeError(f"{attr} not exist")
70

71
    @property
72
    def app(self):
73
        return self._fastapi
74