File size: 1,219 Bytes
ed4d993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import re

from langchain_core.agents 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
        )