Repository: eyurtsev/kor

Issue:
LLMChain is deprecated
Hello,

Seems this project still uses the LLMChain which was deprecated in langchain 0.1.17 and would be removed in langchain 0.3. I do have something that works  I will create a PR for this.

Hints from issue comments:
@uinfante: Bump! Appreciated to get rid of those warnings.

Solution patch:
diff --git a/kor/extraction/api.py b/kor/extraction/api.py
index 5e196c5..8b35d7d 100644
--- a/kor/extraction/api.py
+++ b/kor/extraction/api.py
@@ -1,11 +1,13 @@
 """Kor API for extraction related functionality."""
+
 import asyncio
 from typing import Any, Callable, List, Optional, Sequence, Type, Union, cast
 
-from langchain.chains import LLMChain
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseLanguageModel
+from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import Runnable
 
 from kor.encoders import Encoder, InputFormatter, initialize_encoder
 from kor.extraction.parser import KorParser
@@ -18,7 +20,7 @@ from kor.validators import Validator
 
 async def _extract_from_document_with_semaphore(
     semaphore: asyncio.Semaphore,
-    chain: LLMChain,
+    chain: Runnable,
     document: Document,
     uid: str,
     source_uid: str,
@@ -26,7 +28,7 @@ async def _extract_from_document_with_semaphore(
     """Extract from document with a semaphore to limit concurrency."""
     async with semaphore:
         extraction_result: Extraction = cast(
-            Extraction, await chain.arun(document.page_content)
+            Extraction, await chain.ainvoke(document.page_content)
         )
         return {
             "uid": uid,
@@ -52,7 +54,7 @@ def create_extraction_chain(
     instruction_template: Optional[PromptTemplate] = None,
     verbose: Optional[bool] = None,
     **encoder_kwargs: Any,
-) -> LLMChain:
+) -> Runnable:
     """Create an extraction chain.
     
     Args:
@@ -75,8 +77,10 @@ def create_extraction_chain(
              * "type_description": type description of the node (from TypeDescriptor)
              * "format_instructions": information on how to format the output
                (from Encoder)
-        verbose: if provided, sets the verbosity on the chain, otherwise default
-                 verbosity of the chain will be used
+        verbose: Deprecated, use langchain_core.globals.set_verbose and
+            langchain_core.globals.set_debug instead.
+            Please reference this guide for more information:
+            https://python.langchain.com/v0.2/docs/how_to/debugging
         encoder_kwargs: Keyword arguments to pass to the encoder class
 
     Returns:
@@ -93,32 +97,39 @@ def create_extraction_chain(
         chain = create_extraction_chain(llm, node, encoder_or_encoder_class="JSON",
                                         input_formatter="triple_quotes")
     """
+
+    if verbose is not None:
+        raise NotImplementedError(
+            "The verbose argument is no longer supported. Instead if you want to see "
+            "verbose output, please reference this guide for more information: "
+            "https://python.langchain.com/v0.2/docs/how_to/debugging "
+        )
+
     if not isinstance(node, Object):
         raise ValueError(f"node must be an Object got {type(node)}")
     encoder = initialize_encoder(encoder_or_encoder_class, node, **encoder_kwargs)
     type_descriptor_to_use = initialize_type_descriptors(type_descriptor)
 
-    chain_kwargs = {}
-    if verbose is not None:
-        chain_kwargs["verbose"] = verbose
-
-    return LLMChain(
-        llm=llm,
-        prompt=create_langchain_prompt(
-            node,
-            encoder,
-            type_descriptor_to_use,
-            validator=validator,
-            instruction_template=instruction_template,
-            input_formatter=input_formatter,
-        ),
-        output_parser=KorParser(encoder=encoder, validator=validator, schema_=node),
-        **chain_kwargs,
+    prompt = create_langchain_prompt(
+        node,
+        encoder,
+        type_descriptor_to_use,
+        validator=validator,
+        instruction_template=instruction_template,
+        input_formatter=input_formatter,
+    )
+
+    chain = (
+        prompt
+        | llm
+        | StrOutputParser()
+        | KorParser(encoder=encoder, validator=validator, schema_=node)
     )
+    return chain
 
 
 async def extract_from_documents(
-    chain: LLMChain,
+    chain: Runnable,
     documents: Sequence[Document],
     *,
     max_concurrency: int = 1,


Test patch:
diff --git a/tests/extraction/test_extraction_with_chain.py b/tests/extraction/test_extraction_with_chain.py
index 086ad76..68b3ab8 100644
--- a/tests/extraction/test_extraction_with_chain.py
+++ b/tests/extraction/test_extraction_with_chain.py
@@ -1,10 +1,8 @@
 """Test that the extraction chain works as expected."""
 from typing import Any, Mapping, Optional
 
-import langchain
 import pytest
-from langchain.chains import LLMChain
-from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import Runnable
 
 from kor.encoders import CSVEncoder, JSONEncoder
 from kor.extraction import create_extraction_chain
@@ -40,7 +38,7 @@ def test_create_extraction_chain(options: Mapping[str, Any]) -> None:
 
     for schema in [SIMPLE_OBJECT_SCHEMA]:
         chain = create_extraction_chain(chat_model, schema, **options)
-        assert isinstance(chain, LLMChain)
+        assert isinstance(chain, Runnable)
         # Try to run through predict and parse
         chain.invoke("some string")  # type: ignore
 
@@ -60,7 +58,7 @@ def test_create_extraction_chain_with_csv_encoder(options: Mapping[str, Any]) ->
     chat_model = ToyChatModel(response="hello")
 
     chain = create_extraction_chain(chat_model, **options)
-    assert isinstance(chain, LLMChain)
+    assert isinstance(chain, Runnable)
     # Try to run through predict and parse
     chain.invoke("some string")  # type: ignore
 
@@ -105,46 +103,14 @@ def test_not_implemented_assertion_raised_for_csv(options: Mapping[str, Any]) ->
         create_extraction_chain(chat_model, **options)
 
 
-@pytest.mark.parametrize("verbose", [True, False, None])
+@pytest.mark.parametrize("verbose", [True, False])
 def test_instantiation_with_verbose_flag(verbose: Optional[bool]) -> None:
     """Create an extraction chain."""
     chat_model = ToyChatModel(response="hello")
-    chain = create_extraction_chain(
-        chat_model,
-        SIMPLE_OBJECT_SCHEMA,
-        encoder_or_encoder_class="json",
-        verbose=verbose,
-    )
-    assert isinstance(chain, LLMChain)
-    if verbose is None:
-        expected_verbose = langchain.verbose
-    else:
-        expected_verbose = verbose
-    assert chain.verbose == expected_verbose
-
-
-def test_using_custom_template() -> None:
-    """Create an extraction chain with a custom template."""
-    template = PromptTemplate(
-        input_variables=["format_instructions", "type_description"],
-        template=(
-            "custom_prefix\n"
-            "{type_description}\n\n"
-            "{format_instructions}\n"
-            "custom_suffix"
-        ),
-    )
-    chain = create_extraction_chain(
-        ToyChatModel(response="hello"),
-        OBJECT_SCHEMA_WITH_MANY,
-        instruction_template=template,
-        encoder_or_encoder_class="json",
-    )
-    prompt_value = chain.prompt.format_prompt(text="hello")
-    system_message = prompt_value.to_messages()[0]
-    string_value = prompt_value.to_string()
-
-    assert "custom_prefix" in string_value
-    assert "custom_suffix" in string_value
-    assert "custom_prefix" in system_message.content
-    assert "custom_suffix" in system_message.content
+    with pytest.raises(NotImplementedError):
+        create_extraction_chain(
+            chat_model,
+            SIMPLE_OBJECT_SCHEMA,
+            encoder_or_encoder_class="json",
+            verbose=verbose,
+        )