Callbacks Refactor [base] (#3256)

Co-authored-by: Nuno Campos <nuno@boringbits.io>
Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com>
Co-authored-by: Zander Chase <130414180+vowelparrot@users.noreply.github.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>

tests/integration_tests/callbacks/test_langchain_tracer.py

@@ -0,0 +1,123 @@
+"""Integration tests for the langchain tracer module."""
+import asyncio
+import os
+
+import pytest
+from aiohttp import ClientSession
+
+from langchain.agents import AgentType, initialize_agent, load_tools
+from langchain.callbacks import tracing_enabled
+from langchain.llms import OpenAI
+
+questions = [
+    (
+        "Who won the US Open men's final in 2019? "
+        "What is his age raised to the 0.334 power?"
+    ),
+    (
+        "Who is Olivia Wilde's boyfriend? "
+        "What is his current age raised to the 0.23 power?"
+    ),
+    (
+        "Who won the most recent formula 1 grand prix? "
+        "What is their age raised to the 0.23 power?"
+    ),
+    (
+        "Who won the US Open women's final in 2019? "
+        "What is her age raised to the 0.34 power?"
+    ),
+    ("Who is Beyonce's husband? " "What is his age raised to the 0.19 power?"),
+]
+
+
+def test_tracing_sequential() -> None:
+    os.environ["LANGCHAIN_TRACING"] = "true"
+
+    for q in questions[:3]:
+        llm = OpenAI(temperature=0)
+        tools = load_tools(["llm-math", "serpapi"], llm=llm)
+        agent = initialize_agent(
+            tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+        )
+        agent.run(q)
+
+
+def test_tracing_session_env_var() -> None:
+    os.environ["LANGCHAIN_TRACING"] = "true"
+    os.environ["LANGCHAIN_SESSION"] = "my_session"
+
+    llm = OpenAI(temperature=0)
+    tools = load_tools(["llm-math", "serpapi"], llm=llm)
+    agent = initialize_agent(
+        tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    agent.run(questions[0])
+    if "LANGCHAIN_SESSION" in os.environ:
+        del os.environ["LANGCHAIN_SESSION"]
+
+
+@pytest.mark.asyncio
+async def test_tracing_concurrent() -> None:
+    os.environ["LANGCHAIN_TRACING"] = "true"
+    aiosession = ClientSession()
+    llm = OpenAI(temperature=0)
+    async_tools = load_tools(["llm-math", "serpapi"], llm=llm, aiosession=aiosession)
+    agent = initialize_agent(
+        async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    tasks = [agent.arun(q) for q in questions[:3]]
+    await asyncio.gather(*tasks)
+    await aiosession.close()
+
+
+@pytest.mark.asyncio
+async def test_tracing_concurrent_bw_compat_environ() -> None:
+    os.environ["LANGCHAIN_HANDLER"] = "langchain"
+    if "LANGCHAIN_TRACING" in os.environ:
+        del os.environ["LANGCHAIN_TRACING"]
+    aiosession = ClientSession()
+    llm = OpenAI(temperature=0)
+    async_tools = load_tools(["llm-math", "serpapi"], llm=llm, aiosession=aiosession)
+    agent = initialize_agent(
+        async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    tasks = [agent.arun(q) for q in questions[:3]]
+    await asyncio.gather(*tasks)
+    await aiosession.close()
+    if "LANGCHAIN_HANDLER" in os.environ:
+        del os.environ["LANGCHAIN_HANDLER"]
+
+
+def test_tracing_context_manager() -> None:
+    llm = OpenAI(temperature=0)
+    tools = load_tools(["llm-math", "serpapi"], llm=llm)
+    agent = initialize_agent(
+        tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    if "LANGCHAIN_TRACING" in os.environ:
+        del os.environ["LANGCHAIN_TRACING"]
+    with tracing_enabled() as session:
+        assert session
+        agent.run(questions[0])  # this should be traced
+
+    agent.run(questions[0])  # this should not be traced
+
+
+@pytest.mark.asyncio
+async def test_tracing_context_manager_async() -> None:
+    llm = OpenAI(temperature=0)
+    async_tools = load_tools(["llm-math", "serpapi"], llm=llm)
+    agent = initialize_agent(
+        async_tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    if "LANGCHAIN_TRACING" in os.environ:
+        del os.environ["LANGCHAIN_TRACING"]
+
+    # start a background task
+    task = asyncio.create_task(agent.arun(questions[0]))  # this should not be traced
+    with tracing_enabled() as session:
+        assert session
+        tasks = [agent.arun(q) for q in questions[1:4]]  # these should be traced
+        await asyncio.gather(*tasks)
+
+    await task

tests/integration_tests/callbacks/test_openai_callback.py

@@ -0,0 +1,55 @@
+"""Integration tests for the langchain tracer module."""
+import asyncio
+
+import pytest
+
+from langchain.agents import AgentType, initialize_agent, load_tools
+from langchain.callbacks import get_openai_callback
+from langchain.llms import OpenAI
+
+
+@pytest.mark.asyncio
+async def test_openai_callback() -> None:
+    llm = OpenAI(temperature=0)
+    with get_openai_callback() as cb:
+        llm("What is the square root of 4?")
+
+    total_tokens = cb.total_tokens
+    assert total_tokens > 0
+
+    with get_openai_callback() as cb:
+        llm("What is the square root of 4?")
+        llm("What is the square root of 4?")
+
+    assert cb.total_tokens == total_tokens * 2
+
+    with get_openai_callback() as cb:
+        await asyncio.gather(
+            *[llm.agenerate(["What is the square root of 4?"]) for _ in range(3)]
+        )
+
+    assert cb.total_tokens == total_tokens * 3
+
+    task = asyncio.create_task(llm.agenerate(["What is the square root of 4?"]))
+    with get_openai_callback() as cb:
+        await llm.agenerate(["What is the square root of 4?"])
+
+    await task
+    assert cb.total_tokens == total_tokens
+
+
+def test_openai_callback_agent() -> None:
+    llm = OpenAI(temperature=0)
+    tools = load_tools(["serpapi", "llm-math"], llm=llm)
+    agent = initialize_agent(
+        tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+    )
+    with get_openai_callback() as cb:
+        agent.run(
+            "Who is Olivia Wilde's boyfriend? "
+            "What is his current age raised to the 0.23 power?"
+        )
+        print(f"Total Tokens: {cb.total_tokens}")
+        print(f"Prompt Tokens: {cb.prompt_tokens}")
+        print(f"Completion Tokens: {cb.completion_tokens}")
+        print(f"Total Cost (USD): ${cb.total_cost}")

tests/integration_tests/chains/test_pal.py

@@ -6,7 +6,7 @@ from langchain.chains.pal.base import PALChain
 
 def test_math_prompt() -> None:
     """Test math prompt."""
-    llm = OpenAI(model_name="code-davinci-002", temperature=0, max_tokens=512)
+    llm = OpenAI(temperature=0, max_tokens=512)
     pal_chain = PALChain.from_math_prompt(llm)
     question = (
         "Jan has three times the number of pets as Marcia. "
@@ -19,7 +19,7 @@ def test_math_prompt() -> None:
 
 def test_colored_object_prompt() -> None:
     """Test colored object prompt."""
-    llm = OpenAI(model_name="code-davinci-002", temperature=0, max_tokens=512)
+    llm = OpenAI(temperature=0, max_tokens=512)
     pal_chain = PALChain.from_colored_object_prompt(llm)
     question = (
         "On the desk, you see two blue booklets, "

tests/integration_tests/chains/test_sql_database.py

@@ -27,7 +27,7 @@ def test_sql_database_run() -> None:
     with engine.connect() as conn:
         conn.execute(stmt)
     db = SQLDatabase(engine)
-    db_chain = SQLDatabaseChain(llm=OpenAI(temperature=0), database=db)
+    db_chain = SQLDatabaseChain.from_llm(OpenAI(temperature=0), db)
     output = db_chain.run("What company does Harrison work at?")
     expected_output = " Harrison works at Foo."
     assert output == expected_output
@@ -41,7 +41,7 @@ def test_sql_database_run_update() -> None:
     with engine.connect() as conn:
         conn.execute(stmt)
     db = SQLDatabase(engine)
-    db_chain = SQLDatabaseChain(llm=OpenAI(temperature=0), database=db)
+    db_chain = SQLDatabaseChain.from_llm(OpenAI(temperature=0), db)
     output = db_chain.run("Update Harrison's workplace to Bar")
     expected_output = " Harrison's workplace has been updated to Bar."
     assert output == expected_output
@@ -59,9 +59,7 @@ def test_sql_database_sequential_chain_run() -> None:
     with engine.connect() as conn:
         conn.execute(stmt)
     db = SQLDatabase(engine)
-    db_chain = SQLDatabaseSequentialChain.from_llm(
-        llm=OpenAI(temperature=0), database=db
-    )
+    db_chain = SQLDatabaseSequentialChain.from_llm(OpenAI(temperature=0), db)
     output = db_chain.run("What company does Harrison work at?")
     expected_output = " Harrison works at Foo."
     assert output == expected_output
@@ -77,7 +75,7 @@ def test_sql_database_sequential_chain_intermediate_steps() -> None:
         conn.execute(stmt)
     db = SQLDatabase(engine)
     db_chain = SQLDatabaseSequentialChain.from_llm(
-        llm=OpenAI(temperature=0), database=db, return_intermediate_steps=True
+        OpenAI(temperature=0), db, return_intermediate_steps=True
     )
     output = db_chain("What company does Harrison work at?")
     expected_output = " Harrison works at Foo."

tests/integration_tests/chat_models/test_anthropic.py

@@ -3,7 +3,7 @@ from typing import List
 
 import pytest
 
-from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.chat_models.anthropic import ChatAnthropic
 from langchain.schema import (
     AIMessage,

tests/integration_tests/chat_models/test_openai.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.chat_models.openai import ChatOpenAI
 from langchain.schema import (
     BaseMessage,

tests/integration_tests/chat_models/test_promptlayer_openai.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.chat_models.promptlayer_openai import PromptLayerChatOpenAI
 from langchain.schema import (
     BaseMessage,

tests/integration_tests/llms/test_anthropic.py

@@ -3,7 +3,7 @@ from typing import Generator
 
 import pytest
 
-from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.llms.anthropic import Anthropic
 from langchain.schema import LLMResult
 from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler

tests/integration_tests/llms/test_llamacpp.py

@@ -5,7 +5,6 @@ from typing import Generator
 from urllib.request import urlretrieve
 
 from langchain.llms import LlamaCpp
-from langchain.callbacks.base import CallbackManager
 
 from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
 
@@ -61,10 +60,9 @@ def test_llamacpp_streaming_callback() -> None:
     OFF_BY_ONE = 1  # There may be an off by one error in the upstream code!
 
     callback_handler = FakeCallbackHandler()
-    callback_manager = CallbackManager([callback_handler])
     llm = LlamaCpp(
         model_path=get_model(),
-        callback_manager=callback_manager,
+        callbacks=[callback_handler],
         verbose=True,
         max_tokens=MAX_TOKENS,
     )

tests/integration_tests/llms/test_openai.py

@@ -5,7 +5,7 @@ from typing import Generator
 
 import pytest
 
-from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.llms.loading import load_llm
 from langchain.llms.openai import OpenAI, OpenAIChat
 from langchain.schema import LLMResult

tests/integration_tests/test_schema.py

@@ -1,6 +1,6 @@
 """Test formatting functionality."""
 
-from langchain.schema import _get_num_tokens_default_method
+from langchain.base_language import _get_num_tokens_default_method
 
 
 class TestTokenCountingWithGPT2Tokenizer: