perf: Chat and model

Author: fit2cloud-chenyw

backend/apps/api.py

@@ -3,6 +3,7 @@
 from apps.system.api import login, user, aimodel
 from apps.settings.api import terminology
 from apps.datasource.api import datasource
+from apps.chat.api import chat
 
 
 api_router = APIRouter()
@@ -11,5 +12,6 @@
 api_router.include_router(aimodel.router)
 api_router.include_router(terminology.router)
 api_router.include_router(datasource.router)
+api_router.include_router(chat.router)
 
 

backend/apps/chat/__init__.py

@@ -0,0 +1,53 @@
+from fastapi import APIRouter, HTTPException
+from sqlmodel import select
+from apps.chat.schemas.chat_base_schema import LLMConfig
+from apps.chat.schemas.chat_schema import ChatQuestion
+from apps.chat.schemas.llm import LLMService
+from apps.system.models.system_modle import AiModelDetail
+from common.core.deps import SessionDep
+# from sse_starlette.sse import EventSourceResponse
+router = APIRouter(tags=["Data Q&A"], prefix="/chat")
+
+
+@router.post("/question")
+async def stream_sql(session: SessionDep, requestQuestion: ChatQuestion):
+    question = requestQuestion.question
+    
+    # Use OpenAI model
+    """ openai_config = LLMConfig(
+        model_type="openai",
+        model_name="gpt-4",
+        api_key="your-api-key",
+        additional_params={"temperature": 0.7}
+    )
+    openai_service = LLMService(openai_config) """
+
+    aimodel = session.exec(select(AiModelDetail).where(AiModelDetail.status == True, AiModelDetail.api_key.is_not(None))).first()
+    
+    if not aimodel:
+        raise HTTPException(
+            status_code=400,
+            detail="No available AI model configuration found"
+        )
+    
+    # Use Tongyi Qianwen
+    tongyi_config = LLMConfig(
+        model_type="tongyi",
+        model_name=aimodel.name,
+        api_key=aimodel.api_key,
+        additional_params={"temperature": aimodel.temperature}
+    )
+    llm_service = LLMService(tongyi_config)
+
+    # Use Custom VLLM model
+    """ vllm_config = LLMConfig(
+        model_type="vllm",
+        model_name="your-model-path",
+        additional_params={
+            "max_new_tokens": 200,
+            "temperature": 0.3
+        }
+    )
+    vllm_service = LLMService(vllm_config) """
+    result = llm_service.generate_sql(question)
+    return result

backend/apps/chat/api/__init__.py

@@ -0,0 +1,86 @@
+from pydantic import BaseModel
+from typing import Optional, Dict, Any, Type
+from abc import ABC, abstractmethod
+from langchain_core.language_models import BaseLLM as LangchainBaseLLM
+from langchain_openai import ChatOpenAI
+from langchain_community.llms import Tongyi, VLLM
+
+class LLMConfig(BaseModel):
+    """Base configuration class for large language models"""
+    model_type: str  # Model type: openai/tongyi/vllm etc.
+    model_name: str  # Specific model name
+    api_key: Optional[str] = None
+    api_base_url: Optional[str] = None
+    additional_params: Dict[str, Any] = {}
+    
+
+class BaseLLM(ABC):
+    """Abstract base class for large language models"""
+    
+    def __init__(self, config: LLMConfig):
+        self.config = config
+        self._llm = self._init_llm()
+    
+    @abstractmethod
+    def _init_llm(self) -> LangchainBaseLLM:
+        """Initialize specific large language model instance"""
+        pass
+    
+    @property
+    def llm(self) -> LangchainBaseLLM:
+        """Return the langchain LLM instance"""
+        return self._llm
+
+class OpenAILLM(BaseLLM):
+    def _init_llm(self) -> LangchainBaseLLM:
+        return ChatOpenAI(
+            model=self.config.model_name,
+            api_key=self.config.api_key,
+            **self.config.additional_params
+        )
+    
+    def generate(self, prompt: str) -> str:
+        return self.llm.invoke(prompt)
+
+class TongyiLLM(BaseLLM):
+    def _init_llm(self) -> LangchainBaseLLM:
+        return Tongyi(
+            model_name=self.config.model_name,
+            dashscope_api_key=self.config.api_key,
+            **self.config.additional_params
+        )
+    
+    def generate(self, prompt: str) -> str:
+        return self.llm.invoke(prompt)
+
+class VLLMLLM(BaseLLM):
+    def _init_llm(self) -> LangchainBaseLLM:
+        return VLLM(
+            model=self.config.model_name,
+            **self.config.additional_params
+        )
+    
+    def generate(self, prompt: str) -> str:
+        return self.llm.invoke(prompt)
+
+
+class LLMFactory:
+    """Large Language Model Factory Class"""
+    
+    _llm_types: Dict[str, Type[BaseLLM]] = {
+        "openai": OpenAILLM,
+        "tongyi": TongyiLLM,
+        "vllm": VLLMLLM
+    }
+    
+    @classmethod
+    def create_llm(cls, config: LLMConfig) -> BaseLLM:
+        llm_class = cls._llm_types.get(config.model_type)
+        if not llm_class:
+            raise ValueError(f"Unsupported LLM type: {config.model_type}")
+        return llm_class(config)
+    
+    @classmethod
+    def register_llm(cls, model_type: str, llm_class: Type[BaseLLM]):
+        """Register new model type"""
+        cls._llm_types[model_type] = llm_class

backend/apps/chat/api/chat.py

@@ -0,0 +1,6 @@
+
+from pydantic import BaseModel
+
+
+class ChatQuestion(BaseModel):
+    question: str

backend/apps/chat/schemas/__init__.py

@@ -0,0 +1,33 @@
+from langchain_community.utilities import SQLDatabase
+# from langchain_community.agent_toolkits import create_sql_agent
+from langchain_community.llms import Tongyi
+from langchain_core.prompts import ChatPromptTemplate
+from apps.chat.schemas.chat_base_schema import LLMConfig, LLMFactory
+from common.core.config import settings
+import warnings
+
+warnings.filterwarnings("ignore")
+
+class LLMService:
+    def __init__(self, config: LLMConfig):
+        # Initialize database connection
+        self.db = SQLDatabase.from_uri(str(settings.SQLALCHEMY_DATABASE_URI))
+        
+        # Create LLM instance through factory
+        llm_instance = LLMFactory.create_llm(config)
+        self.llm = llm_instance.llm
+        
+        # Define prompt template
+        self.prompt = ChatPromptTemplate.from_messages([
+            ("system", """You are a professional SQL engineer. Generate PostgreSQL SELECT queries based on the database schema and user questions.
+            Data modification or deletion is prohibited. Table structure is as follows:
+            {schema}
+            """),
+            ("human", "{question}")
+        ])
+    
+    def generate_sql(self, question: str) -> str:
+        chain = self.prompt | self.llm
+        schema = self.db.get_table_info()
+        return chain.invoke({"schema": schema, "question": question})
+

backend/apps/chat/schemas/chat_base_schema.py

@@ -49,7 +49,14 @@ async def dispatch(self, request, call_next):
                 )
             except Exception as e:
                 logging.error(f"Response processing error: {str(e)}", exc_info=True)
-                return response
+                return JSONResponse(
+                    status_code=500,
+                    content={
+                        "code": 500,
+                        "data": None,
+                        "msg": str(e)
+                    }
+                )
 
         return response
 

backend/apps/chat/schemas/chat_schema.py

@@ -3,13 +3,14 @@
 from fastapi.staticfiles import StaticFiles
 import os
 import sentry_sdk
-from fastapi import FastAPI, Path
+from fastapi import FastAPI, Path, HTTPException
 from fastapi.routing import APIRoute
 from starlette.middleware.cors import CORSMiddleware
+from starlette.exceptions import HTTPException as StarletteHTTPException
 from apps.api import api_router
 from apps.system.middleware.auth import TokenMiddleware
 from common.core.config import settings
-from common.core.response_middleware import ResponseMiddleware
+from common.core.response_middleware import ResponseMiddleware, exception_handler
 
 def custom_generate_unique_id(route: APIRoute) -> str:
     tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
@@ -39,6 +40,9 @@ def custom_generate_unique_id(route: APIRoute) -> str:
 app.add_middleware(ResponseMiddleware)
 app.include_router(api_router, prefix=settings.API_V1_STR)
 
+# Register exception handlers
+app.add_exception_handler(StarletteHTTPException, exception_handler.http_exception_handler)
+app.add_exception_handler(Exception, exception_handler.global_exception_handler)
 
 frontend_dist = os.path.abspath("../frontend/dist")
 if not os.path.exists(frontend_dist):