70 lines
2.9 KiB
Python
70 lines
2.9 KiB
Python
|
'''
|
|||
|
|
Author: scutlzc scutlzc@gmail.com
|
|||
|
|
Date: 2024-07-11 16:36:34
|
|||
|
|
LastEditors: scutlzc scutlzc@gmail.com
|
|||
|
|
LastEditTime: 2024-07-12 11:03:24
|
|||
|
|
FilePath: \bizwechat\sqlcode\langchain_model.py
|
|||
|
|
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
|
|||
|
|
'''
|
|||
|
|
from dashscope import Generation
|
|||
|
|
import random
|
|||
|
|
from http import HTTPStatus
|
|||
|
|
from enum import Enum
|
|||
|
|
from .qgi import Generator as BaseGenerator
|
|||
|
|
from langchain_community.llms import Tongyi
|
|||
|
|
from langchain_core.language_models import BaseLLM
|
|||
|
|
from langchain.prompts import PromptTemplate
|
|||
|
|
from langchain.chains.llm import LLMChain
|
|||
|
|
from enum import StrEnum
|
|||
|
|
from .modelloader import ModelManager, ModelLoader
|
|||
|
|
from langchain.agents import AgentExecutor
|
|||
|
|
from langchain_core.exceptions import OutputParserException
|
|||
|
|
from sqlcode.utils import enrich_input
|
|||
|
|
from typing import Optional
|
|||
|
|
from langgraph.graph.graph import CompiledGraph
|
|||
|
|
|
|||
|
|
class Generator(BaseGenerator):
|
|||
|
|
def __init__(self, agentExcutor:CompiledGraph, messages:list[dict[str,str]]|None, apikey:str, seed:int=0, max_retry: Optional[int] = 2) -> None:
|
|||
|
|
self.message = messages
|
|||
|
|
self.apikey = apikey
|
|||
|
|
self.seed = seed
|
|||
|
|
self.agentExcutor = agentExcutor
|
|||
|
|
self.max_retry = max_retry
|
|||
|
|
|
|||
|
|
def _generate(self, input: dict) -> tuple[HTTPStatus]:
|
|||
|
|
seed = self.seed
|
|||
|
|
if seed == 0:
|
|||
|
|
seed = random.randint(1, 10000)
|
|||
|
|
|
|||
|
|
cnt = 0
|
|||
|
|
err = None
|
|||
|
|
while cnt < self.max_retry:
|
|||
|
|
try:
|
|||
|
|
if cnt == 0:
|
|||
|
|
response = self.agentExcutor.invoke(input)["messages"][-1][1]
|
|||
|
|
res = {"input": input, "output": response, "err": None}
|
|||
|
|
else:
|
|||
|
|
new_input = enrich_input("之前agent执行流程发生错误，请模型输出严格按照prompt要求", input)
|
|||
|
|
response = self.agentExcutor.invoke(new_input)["messages"][-1][1]
|
|||
|
|
res = {"input": input, "output": response, "err": err}
|
|||
|
|
status_code = HTTPStatus.OK
|
|||
|
|
except ConnectionError as connectionError:
|
|||
|
|
status_code = HTTPStatus.REQUEST_TIMEOUT
|
|||
|
|
err = str(connectionError)
|
|||
|
|
except OutputParserException as outputException:
|
|||
|
|
cnt+=1
|
|||
|
|
err = str(outputException)
|
|||
|
|
print("outputException: {}".format(outputException))
|
|||
|
|
except KeyError as keyError:
|
|||
|
|
cnt+=1
|
|||
|
|
err = str(keyError)
|
|||
|
|
print("KeyError: {}".format(keyError))
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if cnt == self.max_retry:
|
|||
|
|
status_code = HTTPStatus.INTERNAL_SERVER_ERROR
|
|||
|
|
res = {"input": input, "output": "模型处理异常，建议您将问题表述更加清晰，再重新尝试", "err": err}
|
|||
|
|
|
|||
|
|
return status_code, res
|
|||
|
|
|