extend_test.py 1.69 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
from src.server.extend import LocationExt
from langchain_openai import ChatOpenAI
import re
from src.agent.tool_divisions import complete_administrative_division,divisions



base_llm = ChatOpenAI(
    openai_api_key='xxxxxxxxxxxxx',
    openai_api_base='http://192.168.10.14:8000/v1',
    model_name='Qwen2-7B',
    verbose=True
)

ext = LocationExt(base_llm)
message = [

    ]
question = "大通县和化隆县最大降雨量是多少"
result =  ext.extend_query(question=question,messages=message)
print(result.content)
index = result.content.find("提取到的行政区:")
if index==-1:
    location = result.content
else:
    location = result.content[index + len("提取到的行政区:"):]

# 获取特定子字符串之后的字符串
print(location)
pattern = r'\[([^\]]+)\]'
match = re.search(pattern, location)
cities = []
if match:
    # 匹配到的内容是括号内的全部内容,包括逗号和空格
    # 再次使用split分割这个字符串
    cities = match.group(1).split(', ')
print(cities)
location_exts = []
for m in cities:
    print(m)
    location_ext = complete_administrative_division(m,divisions)
    print(location_ext)
    location_exts.append(location_ext)
prompt1 = ""
locations = []
for l in location_exts:
    if l is not None and "县(区)" in l:
        if isinstance(l["县(区)"], str):
            locations.append(l["县(区)"])
            prompt1 += l["县(区)"] + "位于" + l["省"]+ l["市"]+","
        if isinstance(l["县(区)"], list):
            locations.extend(l["县(区)"])
            prompt1 += l["省"] + l["市"] + "拥有"
            for x in l["县(区)"]:
                prompt1+= x+","

print(locations)
question = prompt1+question
print(question)