Rewrite llm query operator

This is part of the series of tasks for converting llm operator to use
Henry's new LLMMap or LLMElementMap.

Author: bohou-aryn

lib/sycamore/sycamore/docset.py

@@ -1331,10 +1331,9 @@ def llm_query(self, query_agent: LLMTextQueryAgent, **kwargs) -> "DocSet":
             element_type: (Optional) Parameter to only execute the LLM query on a particular element type. If not
                 specified, the query will be executed on all elements.
         """
-        from sycamore.transforms import LLMQuery
+        query = query_agent.as_llm_map(self.plan, **kwargs)
 
-        queries = LLMQuery(self.plan, query_agent=query_agent, **kwargs)
-        return DocSet(self.context, queries)
+        return DocSet(self.context, query)
 
     def groupby(self, grouped_key: Union[str, list[str]], entity: Optional[str] = None) -> "GroupedData":
         from sycamore.grouped_data import GroupedData

lib/sycamore/sycamore/llms/prompts/prompts.py

@@ -662,3 +662,58 @@ def render_element(self, elt: Element, doc: Document) -> RenderedPrompt:
         if self.response_format is not None:
             result.response_format = self.response_format
         return result
+
+
+class JinjaTableMergerPrompt(JinjaElementPrompt):
+    def __init__(
+        self,
+        *,
+        system: Optional[str] = None,
+        user: Union[None, str, list[str]] = None,
+        include_image: bool = False,
+        response_format: ResponseFormat = None,
+        **kwargs,
+    ):
+        from jinja2.sandbox import SandboxedEnvironment
+        from jinja2 import Template
+
+        super().__init__()
+        self.system = system
+        self.user = user
+        self.include_image = include_image
+        self.response_format = response_format
+        self.kwargs = kwargs
+        self._env = SandboxedEnvironment(extensions=["jinja2.ext.loopcontrols"])
+        self._sys_template: Optional[Template] = None
+        self._user_templates: Union[None, list[Template]] = None
+
+    def render_element(self, elt: Element, doc: Document) -> RenderedPrompt:
+        filtered = [e for e in doc.elements if e.type == "table"]
+        idx = filtered.index(elt)
+        prev = None
+        if idx > 0:
+            prev = filtered[idx - 1]
+
+        if self._user_templates is None:
+            userlist = self.user if isinstance(self.user, list) else [self.user]  # type: ignore
+            templates = compile_templates([self.system] + userlist, self._env)  # type: ignore
+            self._sys_template = templates[0]
+            self._user_templates = [t for t in templates[1:] if t is not None]
+
+        render_args = copy.deepcopy(self.kwargs)
+        render_args["elt"] = elt
+        render_args["doc"] = doc
+        render_args["prev"] = prev
+
+        result = render_templates(self._sys_template, self._user_templates, render_args)
+        if self.include_image and len(result.messages) > 0:
+            from sycamore.utils.pdf_utils import get_element_image
+
+            images = []
+            if prev:
+                images.append(get_element_image(prev, doc))
+            images.append(get_element_image(elt, doc))
+            result.messages[-1].images = images
+        if self.response_format is not None:
+            result.response_format = self.response_format
+        return result

lib/sycamore/sycamore/tests/unit/transforms/test_llm_query.py

@@ -15,13 +15,13 @@ def test_llm_query_text_does_not_call_llm(self, mocker):
         prompt = "Give me a one word summary response about the text"
         output_property = "output_property"
         query_agent = LLMTextQueryAgent(prompt=prompt, llm=llm, output_property=output_property)
-        doc = query_agent.execute_query(doc)
+        doc = query_agent.as_llm_map(None)._local_process([doc])
 
-        assert output_property not in doc.elements[0].properties
+        assert output_property not in doc[0].properties
 
     def test_summarize_text_element_calls_llm(self, mocker):
         llm = mocker.Mock(spec=OpenAI)
-        generate = mocker.patch.object(llm, "generate_old")
+        generate = mocker.patch.object(llm, "generate")
         generate.return_value = {"summary": "summary"}
         doc = Document()
         element1 = Element()
@@ -32,14 +32,14 @@ def test_summarize_text_element_calls_llm(self, mocker):
         prompt = "Give me a one word summary response about the text"
         output_property = "output_property"
         query_agent = LLMTextQueryAgent(prompt=prompt, llm=llm, output_property=output_property)
-        doc = query_agent.execute_query(doc)
+        doc = query_agent.as_llm_map(None)._local_process([doc])
 
-        assert doc.elements[0].properties[output_property] == {"summary": "summary"}
-        assert doc.elements[1].properties[output_property] == {"summary": "summary"}
+        assert doc[0].elements[0].properties[output_property] == {"summary": "summary"}
+        assert doc[0].elements[1].properties[output_property] == {"summary": "summary"}
 
     def test_summarize_text_document_calls_llm(self, mocker):
         llm = mocker.Mock(spec=OpenAI)
-        generate = mocker.patch.object(llm, "generate_old")
+        generate = mocker.patch.object(llm, "generate")
         generate.return_value = {"summary": "summary"}
         doc = Document()
         element1 = Element()
@@ -51,6 +51,6 @@ def test_summarize_text_document_calls_llm(self, mocker):
         prompt = "Give me a one word summary response about the text"
         output_property = "output_property"
         query_agent = LLMTextQueryAgent(prompt=prompt, llm=llm, per_element=False, output_property=output_property)
-        doc = query_agent.execute_query(doc)
+        doc = query_agent.as_llm_map(None)._local_process([doc])
 
-        assert doc.properties[output_property] == {"summary": "summary"}
+        assert doc[0].properties[output_property] == {"summary": "summary"}

lib/sycamore/sycamore/transforms/base_llm.py

@@ -188,6 +188,7 @@ def __init__(
         validate: Callable[[Element], bool] = lambda e: True,
         max_tries: int = 5,
         filter: Callable[[Element], bool] = lambda e: True,
+        number_of_elements: Optional[int] = None,
         **kwargs,
     ):
         self._prompt = prompt
@@ -199,6 +200,7 @@ def __init__(
         self._validate = validate
         self._max_tries = max_tries
         self._filter = filter
+        self._number_of_elements = number_of_elements
         super().__init__(child, f=self.llm_map_elements, **kwargs)
 
     def llm_map_elements(self, documents: list[Document]) -> list[Document]:
@@ -207,7 +209,15 @@ def llm_map_elements(self, documents: list[Document]) -> list[Document]:
             for e, _ in elt_doc_pairs:
                 e.properties[self._iteration_var] = 0
 
-        skips = [not self._filter(e) for e, _ in elt_doc_pairs]
+        skips = []
+        counter = 0
+        for e, _ in elt_doc_pairs:
+            if self._filter(e) and (not self._number_of_elements or counter < self._number_of_elements):
+                counter += 1
+                skips.append(False)
+            else:
+                skips.append(True)
+
         tries = 0
         while not all(skips) and tries < self._max_tries:
             tries += 1

lib/sycamore/sycamore/transforms/llm_query.py

@@ -1,8 +1,10 @@
 from typing import Optional, Any, Union
 
 from sycamore.data import Element, Document
+from sycamore.llms.prompts.prompts import JinjaPrompt, ElementListPrompt, JinjaElementPrompt, JinjaTableMergerPrompt
 from sycamore.plan_nodes import NonCPUUser, NonGPUUser, Node
 from sycamore.llms import LLM
+from sycamore.transforms.base_llm import LLMMapElements, LLMMap
 from sycamore.transforms.map import Map
 from sycamore.utils.time_trace import timetrace
 from jinja2.sandbox import SandboxedEnvironment
@@ -55,7 +57,7 @@ def __init__(
         self._output_property = output_property
         self._llm_kwargs = llm_kwargs
         self._per_element = per_element
-        self._format_kwargs = format_kwargs
+        self._format_kwargs = format_kwargs if format_kwargs else {}
         self._number_of_elements = number_of_elements
         self._element_type = element_type
         self._table_cont = table_cont
@@ -125,6 +127,70 @@ def _query_text_object(
                 object["properties"][self._output_property] = llm_resp
         return object
 
+    def get_element_prompt(self):
+        if self._table_cont:
+            prompt = (
+                self._prompt
+                + "\n"
+                + "{% if prev %}"
+                + "ELEMENT 1: \n"
+                + "{prev}"
+                + "\n\n"
+                + "ELEMENT 2: \n"
+                + "{elt}"
+                + "{% else %}"
+                + "{elt}"
+                + "{% endif %}"
+            )
+            return JinjaTableMergerPrompt(system=None, user=prompt, **self._format_kwargs)
+        else:
+            return JinjaElementPrompt(system=None, user=self._prompt, **self._format_kwargs)
+
+    def get_document_prompt(self):
+        if self._number_of_elements:
+
+            def selector(elts: list[Element]):
+                elts = [elt for elt in elts if elt.type == self._element_type]
+                if self._number_of_elements:
+                    elts = elts[: self._number_of_elements]
+                return elts
+
+            return ElementListPrompt(
+                system=None,
+                user=self._prompt,
+                element_select=selector,
+                element_list_constructor=lambda elts: "\n".join([elt.text_representation for elt in elts]),
+                **self._format_kwargs,
+            )
+        else:
+            if self._format_kwargs:
+                return JinjaPrompt(system=None, user=self._prompt, **self._format_kwargs)
+            else:
+                return JinjaPrompt(system=None, user=self._prompt + "\n" + "{{doc}}", **self._format_kwargs)
+
+    def as_llm_map(self, child) -> Node:
+        if self._per_element:
+
+            def element_filter(e: Element):
+                return not self._element_type or e.type == self._element_type
+
+            prompt = self.get_element_prompt()
+            llm_map = LLMMapElements(
+                child,
+                prompt=prompt,
+                output_field=self._output_property,
+                llm=self._llm,
+                filter=element_filter,
+                number_of_elements=self._number_of_elements,
+                **self._llm_kwargs,
+            )
+        else:
+            prompt = self.get_document_prompt()
+            llm_map = LLMMap(
+                child, prompt=prompt, output_field=self._output_property, llm=self._llm, **self._llm_kwargs
+            )
+        return llm_map
+
 
 class LLMQuery(NonCPUUser, NonGPUUser, Map):
     """