docs: put documentation in the docs folder instead of examples (#1674)

docs/docs/tutorials/llm-compiler/math_tools.py

@@ -0,0 +1,142 @@
+import math
+import re
+from typing import List, Optional
+
+import numexpr
+from langchain.chains.openai_functions import create_structured_output_runnable
+from langchain_core.messages import SystemMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import StructuredTool
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel, Field
+
+_MATH_DESCRIPTION = (
+    "math(problem: str, context: Optional[list[str]]) -> float:\n"
+    " - Solves the provided math problem.\n"
+    ' - `problem` can be either a simple math problem (e.g. "1 + 3") or a word problem (e.g. "how many apples are there if there are 3 apples and 2 apples").\n'
+    " - You cannot calculate multiple expressions in one call. For instance, `math('1 + 3, 2 + 4')` does not work. "
+    "If you need to calculate multiple expressions, you need to call them separately like `math('1 + 3')` and then `math('2 + 4')`\n"
+    " - Minimize the number of `math` actions as much as possible. For instance, instead of calling "
+    '2. math("what is the 10% of $1") and then call 3. math("$1 + $2"), '
+    'you MUST call 2. math("what is the 110% of $1") instead, which will reduce the number of math actions.\n'
+    # Context specific rules below
+    " - You can optionally provide a list of strings as `context` to help the agent solve the problem. "
+    "If there are multiple contexts you need to answer the question, you can provide them as a list of strings.\n"
+    " - `math` action will not see the output of the previous actions unless you provide it as `context`. "
+    "You MUST provide the output of the previous actions as `context` if you need to do math on it.\n"
+    " - You MUST NEVER provide `search` type action's outputs as a variable in the `problem` argument. "
+    "This is because `search` returns a text blob that contains the information about the entity, not a number or value. "
+    "Therefore, when you need to provide an output of `search` action, you MUST provide it as a `context` argument to `math` action. "
+    'For example, 1. search("Barack Obama") and then 2. math("age of $1") is NEVER allowed. '
+    'Use 2. math("age of Barack Obama", context=["$1"]) instead.\n'
+    " - When you ask a question about `context`, specify the units. "
+    'For instance, "what is xx in height?" or "what is xx in millions?" instead of "what is xx?"\n'
+)
+
+
+_SYSTEM_PROMPT = """Translate a math problem into a expression that can be executed using Python's numexpr library. Use the output of running this code to answer the question.
+
+Question: ${{Question with math problem.}}
+```text
+${{single line mathematical expression that solves the problem}}
+```
+...numexpr.evaluate(text)...
+```output
+${{Output of running the code}}
+```
+Answer: ${{Answer}}
+
+Begin.
+
+Question: What is 37593 * 67?
+ExecuteCode({{code: "37593 * 67"}})
+...numexpr.evaluate("37593 * 67")...
+```output
+2518731
+```
+Answer: 2518731
+
+Question: 37593^(1/5)
+ExecuteCode({{code: "37593**(1/5)"}})
+...numexpr.evaluate("37593**(1/5)")...
+```output
+8.222831614237718
+```
+Answer: 8.222831614237718
+"""
+
+_ADDITIONAL_CONTEXT_PROMPT = """The following additional context is provided from other functions.\
+    Use it to substitute into any ${{#}} variables or other words in the problem.\
+    \n\n${context}\n\nNote that context variables are not defined in code yet.\
+You must extract the relevant numbers and directly put them in code."""
+
+
+class ExecuteCode(BaseModel):
+    """The input to the numexpr.evaluate() function."""
+
+    reasoning: str = Field(
+        ...,
+        description="The reasoning behind the code expression, including how context is included, if applicable.",
+    )
+
+    code: str = Field(
+        ...,
+        description="The simple code expression to execute by numexpr.evaluate().",
+    )
+
+
+def _evaluate_expression(expression: str) -> str:
+    try:
+        local_dict = {"pi": math.pi, "e": math.e}
+        output = str(
+            numexpr.evaluate(
+                expression.strip(),
+                global_dict={},  # restrict access to globals
+                local_dict=local_dict,  # add common mathematical functions
+            )
+        )
+    except Exception as e:
+        raise ValueError(
+            f'Failed to evaluate "{expression}". Raised error: {repr(e)}.'
+            " Please try again with a valid numerical expression"
+        )
+
+    # Remove any leading and trailing brackets from the output
+    return re.sub(r"^\[|\]$", "", output)
+
+
+def get_math_tool(llm: ChatOpenAI):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", _SYSTEM_PROMPT),
+            ("user", "{problem}"),
+            MessagesPlaceholder(variable_name="context", optional=True),
+        ]
+    )
+    extractor = prompt | llm.with_structured_output(ExecuteCode)
+
+    def calculate_expression(
+        problem: str,
+        context: Optional[List[str]] = None,
+        config: Optional[RunnableConfig] = None,
+    ):
+        chain_input = {"problem": problem}
+        if context:
+            context_str = "\n".join(context)
+            if context_str.strip():
+                context_str = _ADDITIONAL_CONTEXT_PROMPT.format(
+                    context=context_str.strip()
+                )
+                chain_input["context"] = [SystemMessage(content=context_str)]
+        code_model = extractor.invoke(chain_input, config)
+        try:
+            return _evaluate_expression(code_model.code)
+        except Exception as e:
+            return repr(e)
+
+    return StructuredTool.from_function(
+        name="math",
+        func=calculate_expression,
+        description=_MATH_DESCRIPTION,
+    )

docs/docs/tutorials/llm-compiler/output_parser.py

@@ -0,0 +1,177 @@
+import ast
+import re
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+from langchain_core.exceptions import OutputParserException
+from langchain_core.messages import BaseMessage
+from langchain_core.output_parsers.transform import BaseTransformOutputParser
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from typing_extensions import TypedDict
+
+THOUGHT_PATTERN = r"Thought: ([^\n]*)"
+ACTION_PATTERN = r"\n*(\d+)\. (\w+)\((.*)\)(\s*#\w+\n)?"
+# $1 or ${1} -> 1
+ID_PATTERN = r"\$\{?(\d+)\}?"
+END_OF_PLAN = "<END_OF_PLAN>"
+
+
+### Helper functions
+
+
+def _ast_parse(arg: str) -> Any:
+    try:
+        return ast.literal_eval(arg)
+    except:  # noqa
+        return arg
+
+
+def _parse_llm_compiler_action_args(args: str, tool: Union[str, BaseTool]) -> list[Any]:
+    """Parse arguments from a string."""
+    if args == "":
+        return ()
+    if isinstance(tool, str):
+        return ()
+    extracted_args = {}
+    tool_key = None
+    prev_idx = None
+    for key in tool.args.keys():
+        # Split if present
+        if f"{key}=" in args:
+            idx = args.index(f"{key}=")
+            if prev_idx is not None:
+                extracted_args[tool_key] = _ast_parse(
+                    args[prev_idx:idx].strip().rstrip(",")
+                )
+            args = args.split(f"{key}=", 1)[1]
+            tool_key = key
+            prev_idx = 0
+    if prev_idx is not None:
+        extracted_args[tool_key] = _ast_parse(
+            args[prev_idx:].strip().rstrip(",").rstrip(")")
+        )
+    return extracted_args
+
+
+def default_dependency_rule(idx, args: str):
+    matches = re.findall(ID_PATTERN, args)
+    numbers = [int(match) for match in matches]
+    return idx in numbers
+
+
+def _get_dependencies_from_graph(
+    idx: int, tool_name: str, args: Dict[str, Any]
+) -> dict[str, list[str]]:
+    """Get dependencies from a graph."""
+    if tool_name == "join":
+        return list(range(1, idx))
+    return [i for i in range(1, idx) if default_dependency_rule(i, str(args))]
+
+
+class Task(TypedDict):
+    idx: int
+    tool: BaseTool
+    args: list
+    dependencies: Dict[str, list]
+    thought: Optional[str]
+
+
+def instantiate_task(
+    tools: Sequence[BaseTool],
+    idx: int,
+    tool_name: str,
+    args: Union[str, Any],
+    thought: Optional[str] = None,
+) -> Task:
+    if tool_name == "join":
+        tool = "join"
+    else:
+        try:
+            tool = tools[[tool.name for tool in tools].index(tool_name)]
+        except ValueError as e:
+            raise OutputParserException(f"Tool {tool_name} not found.") from e
+    tool_args = _parse_llm_compiler_action_args(args, tool)
+    dependencies = _get_dependencies_from_graph(idx, tool_name, tool_args)
+
+    return Task(
+        idx=idx,
+        tool=tool,
+        args=tool_args,
+        dependencies=dependencies,
+        thought=thought,
+    )
+
+
+class LLMCompilerPlanParser(BaseTransformOutputParser[dict], extra="allow"):
+    """Planning output parser."""
+
+    tools: List[BaseTool]
+
+    def _transform(self, input: Iterator[Union[str, BaseMessage]]) -> Iterator[Task]:
+        texts = []
+        # TODO: Cleanup tuple state tracking here.
+        thought = None
+        for chunk in input:
+            # Assume input is str. TODO: support vision/other formats
+            text = chunk if isinstance(chunk, str) else str(chunk.content)
+            for task, thought in self.ingest_token(text, texts, thought):
+                yield task
+        # Final possible task
+        if texts:
+            task, _ = self._parse_task("".join(texts), thought)
+            if task:
+                yield task
+
+    def parse(self, text: str) -> List[Task]:
+        return list(self._transform([text]))
+
+    def stream(
+        self,
+        input: str | BaseMessage,
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> Iterator[Task]:
+        yield from self.transform([input], config, **kwargs)
+
+    def ingest_token(
+        self, token: str, buffer: List[str], thought: Optional[str]
+    ) -> Iterator[Tuple[Optional[Task], str]]:
+        buffer.append(token)
+        if "\n" in token:
+            buffer_ = "".join(buffer).split("\n")
+            suffix = buffer_[-1]
+            for line in buffer_[:-1]:
+                task, thought = self._parse_task(line, thought)
+                if task:
+                    yield task, thought
+            buffer.clear()
+            buffer.append(suffix)
+
+    def _parse_task(self, line: str, thought: Optional[str] = None):
+        task = None
+        if match := re.match(THOUGHT_PATTERN, line):
+            # Optionally, action can be preceded by a thought
+            thought = match.group(1)
+        elif match := re.match(ACTION_PATTERN, line):
+            # if action is parsed, return the task, and clear the buffer
+            idx, tool_name, args, _ = match.groups()
+            idx = int(idx)
+            task = instantiate_task(
+                tools=self.tools,
+                idx=idx,
+                tool_name=tool_name,
+                args=args,
+                thought=thought,
+            )
+            thought = None
+        # Else it is just dropped
+        return task, thought