Spaces:
Running
Running
| import time | |
| from fastapi import APIRouter, HTTPException, status | |
| from fastapi import HTTPException, status | |
| from reggpt.schemas.schema import UserQuery, LoginRequest, UserModel | |
| from reggpt.routers.controller import get_QA_Answers, get_avaliable_models | |
| from reggpt.configs.api import API_ENDPOINT_LOGIN,API_ENDPOINT_CHAT, API_ENDPOINT_HEALTH, API_ENDPOINT_MODEL | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class ChatAPI: | |
| def __init__(self): | |
| self.router = APIRouter() | |
| self.router.add_api_route(API_ENDPOINT_HEALTH, self.hello, methods=["GET"]) | |
| self.router.add_api_route(API_ENDPOINT_MODEL, self.avaliable_models, methods=["GET"]) | |
| self.router.add_api_route( | |
| API_ENDPOINT_LOGIN, self.login, methods=["POST"], response_model=UserModel | |
| ) | |
| self.router.add_api_route(API_ENDPOINT_CHAT, self.chat, methods=["POST"]) | |
| async def hello(self): | |
| return "Hello there!" | |
| async def avaliable_models(self): | |
| logger.info("getting avaliable models") | |
| models = get_avaliable_models() | |
| if not models: | |
| logger.exception("models not found") | |
| raise HTTPException( | |
| status_code=status.HTTP_404_NOT_FOUND, detail="models not found" | |
| ) | |
| return models | |
| async def login(self, login_request: LoginRequest): | |
| logger.info(f"username password: {login_request} ") | |
| # Dummy user data for demonstration (normally, you'd use a database) | |
| dummy_users_db = { | |
| "john_doe": { | |
| "userId": 1, | |
| "firstName": "John", | |
| "lastName": "Doe", | |
| "userName": "john_doe", | |
| "password": "password", # Normally, passwords would be hashed and stored securely | |
| "token": "dummy_token_123", # In a real scenario, this would be a JWT or another kind of token | |
| } | |
| } | |
| # Fetch user by username | |
| # user = dummy_users_db.get(login_request.username) | |
| user = dummy_users_db.get("john_doe") | |
| # Validate user credentials | |
| if not user or user["password"] != login_request.password: | |
| raise HTTPException(status_code=401, detail="Invalid username or password") | |
| # Return the user model without the password | |
| return UserModel( | |
| userId=user["userId"], | |
| firstName=user["firstName"], | |
| lastName=user["lastName"], | |
| userName=user["userName"], | |
| token=user["token"], | |
| ) | |
| async def chat( | |
| self, userQuery: UserQuery | |
| ): #:UserQuery):# -> ResponseModel: #chat: QueryModel): # -> ResponseModel: | |
| """Makes query to doc store via Langchain pipeline. | |
| :param chat.: question, model, dataset location, history of the chat. | |
| :type chat: QueryModel | |
| """ | |
| logger.info(f"userQuery: {userQuery} ") | |
| try: | |
| start = time.time() | |
| res = get_QA_Answers(userQuery) | |
| logger.info( | |
| f"-------------------------- answer: {res} -------------------------- " | |
| ) | |
| # return res | |
| end = time.time() | |
| logger.info( | |
| f"-------------------------- Server process (took {round(end - start, 2)} s.) \n: {res}" | |
| ) | |
| print( | |
| f" \n -------------------------- Server process (took {round(end - start, 2)} s.) ------------------------- \n" | |
| ) | |
| # return res | |
| return { | |
| "user_input": userQuery.user_question, | |
| "bot_response": res["bot_response"], | |
| "format_data": res["format_data"], | |
| } | |
| except HTTPException as e: | |
| logger.exception(e) | |
| raise e | |
| except Exception as e: | |
| logger.exception(e) | |
| raise HTTPException(status_code=400, detail=f"Error : {e}") | |