mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
import re
|
|
|
|
from langchain.schema.agent import AgentAction, AgentFinish
|
|
|
|
from .agent_scratchpad import _format_docs
|
|
|
|
|
|
def extract_between_tags(tag: str, string: str, strip: bool = True) -> str:
|
|
ext_list = re.findall(f"<{tag}\s?>(.+?)</{tag}\s?>", string, re.DOTALL)
|
|
if strip:
|
|
ext_list = [e.strip() for e in ext_list]
|
|
if ext_list:
|
|
if len(ext_list) != 1:
|
|
raise ValueError
|
|
# Only return the first one
|
|
return ext_list[0]
|
|
|
|
|
|
def parse_output(outputs):
|
|
partial_completion = outputs["partial_completion"]
|
|
steps = outputs["intermediate_steps"]
|
|
search_query = extract_between_tags(
|
|
"search_query", partial_completion + "</search_query>"
|
|
)
|
|
if search_query is None:
|
|
docs = []
|
|
str_output = ""
|
|
for action, observation in steps:
|
|
docs.extend(observation)
|
|
str_output += action.log
|
|
str_output += "</search_query>" + _format_docs(observation)
|
|
str_output += partial_completion
|
|
return AgentFinish({"docs": docs, "output": str_output}, log=partial_completion)
|
|
else:
|
|
return AgentAction(
|
|
tool="search", tool_input=search_query, log=partial_completion
|
|
)
|