41 lines
1.6 KiB
Python
41 lines
1.6 KiB
Python
import sys
|
||
import os
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||
from configs import LLM_MODELS, TEMPERATURE
|
||
from server.utils import get_ChatOpenAI
|
||
from langchain.chains import LLMChain
|
||
from langchain.agents import LLMSingleActionAgent, AgentExecutor
|
||
from server import tools, tool_names
|
||
from langchain.memory import ConversationBufferWindowMemory
|
||
|
||
memory = ConversationBufferWindowMemory(k=5)
|
||
model = get_ChatOpenAI(
|
||
model_name=LLM_MODELS[0],
|
||
temperature=TEMPERATURE,
|
||
)
|
||
from server.agent.custom_template import CustomOutputParser, prompt
|
||
|
||
output_parser = CustomOutputParser()
|
||
llm_chain = LLMChain(llm=model, prompt=prompt)
|
||
agent = LLMSingleActionAgent(
|
||
llm_chain=llm_chain,
|
||
output_parser=output_parser,
|
||
stop=["\nObservation:"],
|
||
allowed_tools=tool_names
|
||
)
|
||
|
||
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, memory=memory, verbose=True)
|
||
|
||
import pytest
|
||
@pytest.mark.parametrize("text_prompt",
|
||
["北京市朝阳区未来24小时天气如何？", # 天气功能函数
|
||
"计算 (2 + 2312312)/4 是多少？", # 计算功能函数
|
||
"翻译这句话成中文：Life is the art of drawing sufficient conclusions form insufficient premises."] # 翻译功能函数
|
||
)
|
||
def test_different_agent_function(text_prompt):
|
||
try:
|
||
text_answer = agent_executor.run(text_prompt)
|
||
assert text_answer is not None
|
||
except Exception as e:
|
||
pytest.fail(f"agent_function failed with {text_prompt}, error: {str(e)}")
|