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)