---

Continue Exploring Chat Models with LangChain

Data Mastery Series — Episode 31: LangChain Website (Part 6)

Connect with me and follow our journey: Linkedin, Facebook

---

Welcome to Episode 31 of the Data Mastery Series, where we continue our deep dive into LangChain’s powerful tools and capabilities. In the previous episodes, we covered Model I/O, Prompt Templates, Few-Shot Prompts, and Chat Models. You can catch up here:

• Part 1: LangChain Model I/O Basics
• Part 2: Unpacking Prompt Templates with LangChain
• Part 3: Exploring Few-Shot Prompts with LangChain
• Part 4: Unpacking Chat Models with LangChain
• Part 5: Exploring Chat Models with LangChain

Note: As always, I’ll be sharing insights and practical examples based on my own experience with the LangChain documentation. Let’s jump in!

Source: https://python.langchain.com/v0.1/docs/modules/model%5Fio/chat/

In this episode, we’ll explore how LangChain enhances structured outputs, caching, custom chat models, log probabilities, response metadata, and token usage tracking. Whether you’re a seasoned developer or just starting with LLMs, these features can make your projects more efficient, scalable, and cost-effective. Let’s dive in!

1. Structured Output: Why Does It Matter?

Getting structured data (like JSON or XML) back from LLMs is crucial for using them in applications. Here’s how LangChain helps:

• Prompting: Simply ask the LLM to format its response. This is easy but doesn’t guarantee perfect formatting.
• Function Calling: The LLM can return a function call based on the provided schema, ensuring well-structured results.
• Tool Calling: Like function calling, but for multiple functions.
• JSON Mode: The LLM returns data in JSON format, ensuring structure.

LangChain’s .with_structured_output simplifies this process. Let’s see an example using Pydantic and OpenAI:

Example - Using Pydantic to define the output structure with OpenAI:

class Joke(BaseModel):
setup: str = Field(description="The setup of the joke")
punchline: str = Field(description="The punchline to the joke")

Uses Function Calling by default

model = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, api_key = OPENAI_API_KEY)
structured_llm = model.with_structured_output(Joke)

Invoke the model with a query

structured_llm.invoke("Tell me a joke about cats")

Output:

'''
Joke(setup='Why was the cat sitting on the computer?', punchline='To keep an eye on the mouse!')
'''

Example - Using JSON Mode:

structured_llm = model.with_structured_output(Joke, method="json_mode")
result = structured_llm.invoke("Tell me a joke about cats, respond in JSON with setup and punchline keys")
print(result.json())

Output:

'''
{
"setup": "Why don't cats play poker in the jungle?",
"punchline": "Too many cheetahs!"
}
'''

2. Caching: Saving Time and Money

Caching previous LLM responses can dramatically improve efficiency:

• Cost Savings: Reduce API calls, especially for repeated queries.
• Speed Boost: Retrieve cached responses instantly.

Example - Caching (First call takes longer as it's not cached)

llm = ChatOpenAI(model="gpt-3.5-turbo-0125", api_key = OPENAI_API_KEY)

Set up in-memory cache

set_llm_cache(InMemoryCache())

%%time
llm.predict("Tell me a joke")

Output of 1st run:

'''
CPU times: user 50.5 ms, sys: 2.71 ms, total: 53.2 ms
Wall time: 893 ms
Why couldn't the bicycle stand up by itself?\n\nBecause it was two-tired!
'''

Example - Caching (Second call is faster as it retrieves from cache)

%%time
llm.predict("Tell me a joke")

Output:

'''
CPU times: user 918 µs, sys: 0 ns, total: 918 µs
Wall time: 2.44 ms
Why couldn't the bicycle stand up by itself?\n\nBecause it was two-tired!
'''

3. Custom Chat Models: Tailoring Your AI

At its core, a chat model takes input in the form of text and returns a response, also in text format. LangChain makes it easy to work with this structure by providing pre-defined message types. These include:

• SystemMessage: Sets the AI’s behavior (e.g., helpful assistant, expert).
• HumanMessage: Represents messages from the user.
• AIMessage: Represents the AI’s response.
• FunctionMessage/ToolMessage: Sends tool results back to the model.
• AIMessageChunk/HumanMessageChunk: Smaller pieces of messages for streaming data from the chat model.

A) Streaming:

Example - Combining message chunks to create a full response

AIMessageChunk(content="Hello") + AIMessageChunk(content=" World!")

Output:

'''
AIMessageChunk(content='Hello World!', additional_kwargs={}, response_metadata={})
'''

B) Base Chat Model:

In LangChain, a Base Chat Model requires you to implement key methods for handling chat input and output. Here’s a concise overview of these methods:

• _generate: Generates chat results from the prompt. (Required)
• _llm_type: Specifies the model type for logging. (Required)
• _identifying_params: Shows the model’s parameters for tracking. (Optional)
• _stream: For streaming output. (Optional)
• _agenerate: Async version of _generate. (Optional)
• _astream: Async version of _stream. (Optional)

C) Implementation:

Let’s dive into a practical example of a custom chat model that returns the first n characters of the last message in a prompt. This model demonstrates the flexibility of LangChain for creating tailored AI models.

Example - Base Chat Model (Multiple Messages)

class CustomChatModelAdvanced(BaseChatModel):
"""A custom chat model that echoes the first n characters of the input."""

model_name: str  
n: int  

def _generate(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> ChatResult:  
    last_message = messages[-1]  
    tokens = last_message.content[: self.n]  
    message = AIMessage(  
        content=tokens,  
        additional_kwargs={},  # Used to add additional payload (e.g., function calling request)  
        response_metadata={  # Use for response metadata  
            "time_in_seconds": 3,  
        },  
    )  

    generation = ChatGeneration(message=message)  
    return ChatResult(generations=[generation])  

def _stream(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> Iterator[ChatGenerationChunk]:  
    last_message = messages[-1]  
    tokens = last_message.content[: self.n]  

    for token in tokens:  
        chunk = ChatGenerationChunk(message=AIMessageChunk(content=token))  
        if run_manager:  
            run_manager.on_llm_new_token(token, chunk=chunk)  

        yield chunk  

    # Let's add some other information (e.g., response metadata)  
    chunk = ChatGenerationChunk(  
        message=AIMessageChunk(content="", response_metadata={"time_in_sec": 3})  
    )  
    if run_manager:  
        run_manager.on_llm_new_token(token, chunk=chunk)  
    yield chunk  

@property  
def _llm_type(self) -> str:  
    return "echoing-chat-model-advanced"  

@property  
def _identifying_params(self) -> Dict[str, Any]:  
    return {"model_name": self.model_name}  

Testing the Custom Model

model = CustomChatModelAdvanced(n=3, model_name="my_custom_model")
model.invoke(
[
HumanMessage(content="hello!"),
AIMessage(content="Hi there human!"),
HumanMessage(content="Meow!"),
]
)

Output:

'''
AIMessage(content='Meo', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-673e1dc6-2a2b-4775-ae7e-df2c3a3ca02c-0')
'''

Example - Base Chat Model (Single Message)

model.invoke("hello")

Output:

'''
AIMessage(content='hel', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-8813fc24-9733-4eb1-bbd9-325dcff91f8c-0')
'''

Example - Base Chat Model (Batch Processing)

model.batch(["hello", "goodbye"])

Output:

'''
[AIMessage(content='hel', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-8813fc24-9733-4eb1-bbd9-325dcff91f8c-0'),
AIMessage(content='goo', additional_kwargs={}, response_metadata={'time_in_seconds': 3}, id='run-952a67ee-f11c-4fdf-9b5c-eb89b4d35bd1-0')]
'''

Example - Base Chat Model (Streaming Data)

for chunk in model.stream("cat"):
print(chunk.content, end="|")

Output:

'''
c|a|t||
'''

Example - Base Chat Model (Asynchronous Streaming)

async for chunk in model.astream("cat"):
print(chunk.content, end="|")

Output:

'''
c|a|t||
'''

Example - Base Chat Model (Using API astream_events)

async for event in model.astream_events("cat", version="v1"):
print(event)

Output:

'''
'event': 'on_chat_model_start', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'name': 'CustomChatModelAdvanced', 'tags': [], 'metadata': {}, 'data': {'input': 'cat'}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='c', additional_kwargs={}, response_metadata={}, id='run-7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='a', additional_kwargs={}, response_metadata={}, id='run-7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='t', additional_kwargs={}, response_metadata={}, id='run-7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c')}, 'parent_ids': []}
{'event': 'on_chat_model_stream', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c')}, 'parent_ids': []}
{'event': 'on_chat_model_end', 'name': 'CustomChatModelAdvanced', 'run_id': '7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c', 'tags': [], 'metadata': {}, 'data': {'output': AIMessageChunk(content='cat', additional_kwargs={}, response_metadata={'time_in_sec': 3}, id='run-7714fb10-3a3b-4f9b-a2ae-3f8ea3c5e49c')}, 'parent_ids': []}
'''

D) Contributing (Summary)

When contributing a chat model integration to LangChain, make sure to include:

• Documentation: Provide clear docstrings for all arguments, along with links to the model’s API if applicable.
• Tests: Include unit or integration tests for each overridden method, especially for invoke, stream, and batch.
• Streaming (Optional): Implement the _stream method.
• Stop Token Behavior: Ensure that stop tokens are respected and included in responses.
• Secret API Keys: Use SecretStr for any API keys to prevent them from being exposed accidentally.
• Identifying Params: Include a model_name
• Optimizations: Consider implementing asynchronous versions (_agenerate and _astream) for better performance.

4. Get log probabilities: Understanding AI Decisions

Some OpenAI models can return log probabilities for each token in the response, which helps explain why the model selected certain words. To retrieve log probabilities, set logprobs=True when calling the API.

Example - Retrieving Log Probabilities:

llm = ChatOpenAI(model="gpt-3.5-turbo-0125", api_key = OPENAI_API_KEY).bind(logprobs=True)
msg = llm.invoke("human", "how are you today"))
print(msg.response_metadata["logprobs"]["content"][:5]) # Accessing log probabilities

Output:

'''
[{'token': 'I', 'bytes': [73], 'logprob': -0.27539697, 'top_logprobs': []},
{'token': "'m", 'bytes': [39, 109], 'logprob': -0.39702156, 'top_logprobs': []},
{'token': ' just', 'bytes': [32, 106, 117, 115, 116], 'logprob': -0.23634185, 'top_logprobs': []},
{'token': ' a', 'bytes': [32, 97], 'logprob': -0.0020874506, 'top_logprobs': []},
{'token': ' computer', 'bytes': [32, 99, 111, 109, 112, 117, 116, 101, 114], 'logprob': -0.054618873, 'top_logprobs': []}]
'''

Example - Streaming Log Probabilities:

ct = 0
full = None
for chunk in llm.stream(("human", "how are you today")):
if ct < 5:
full = chunk if full is None else full + chunk
if "logprobs" in full.response_metadata:
print(full.response_metadata["logprobs"]["content"])
else:
break
ct += 1

Output:

'''
[]
[{'token': 'I', 'bytes': [73], 'logprob': -0.27539697, 'top_logprobs': []}]
[{'token': 'I', 'bytes': [73], 'logprob': -0.27539697, 'top_logprobs': []}, {'token': "'m", 'bytes': [39, 109], 'logprob': -0.39702156, 'top_logprobs': []}]
[{'token': 'I', 'bytes': [73], 'logprob': -0.27539697, 'top_logprobs': []}, {'token': "'m", 'bytes': [39, 109], 'logprob': -0.39702156, 'top_logprobs': []}, {'token': ' just', 'bytes': [32, 106, 117, 115, 116], 'logprob': -0.23634185, 'top_logprobs': []}]
[{'token': 'I', 'bytes': [73], 'logprob': -0.27539697, 'top_logprobs': []}, {'token': "'m", 'bytes': [39, 109], 'logprob': -0.39702156, 'top_logprobs': []}, {'token': ' just', 'bytes': [32, 106, 117, 115, 116], 'logprob': -0.23634185, 'top_logprobs': []}, {'token': ' a', 'bytes': [32, 97], 'logprob': -0.0020874506, 'top_logprobs': []}]
'''

5. Response Metadata: Insights into Model Behavior

Response metadata from language models provides information about the generated output. This can include token usage, time taken, and log probabilities. This data helps analyze and improve model usage.

Example - Response Metadata

llm = ChatOpenAI(model="gpt-4-turbo", api_key = OPENAI_API_KEY)
msg = llm.invoke([("human", "What's the oldest known example of cuneiform")])
print(msg.response_metadata)

Output:

'''
{'token_usage': {'completion_tokens': 215,
'prompt_tokens': 17,
'total_tokens': 232,
'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0},
'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}},
'model_name': 'gpt-4-turbo-2024-04-09',
'system_fingerprint': 'fp_83975a045a',
'finish_reason': 'stop',
'logprobs': None}
'''

6. Tracking Token Usage: Manage Costs and Performance

Every call to an LLM uses tokens, and tracking these tokens helps you understand the cost and optimize performance. You can track token usage either by reviewing response metadata or by using LangChain’s callbacks.

Example - Tracking Token Usage with response_metadata:

llm = ChatOpenAI(model="gpt-4-turbo", api_key = OPENAI_API_KEY)
msg = llm.invoke([("human", "What's the oldest known example of cuneiform")])
print(msg.response_metadata)

Output:

'''
{'token_usage': {'completion_tokens': 215,
'prompt_tokens': 17,
'total_tokens': 232,
'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0},
'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}},
'model_name': 'gpt-4-turbo-2024-04-09',
'system_fingerprint': 'fp_83975a045a',
'finish_reason': 'stop',
'logprobs': None}
'''

Example - Tracking Token Usage with Callbacks:

llm = ChatOpenAI(model="gpt-4-turbo", temperature=0, api_key = OPENAI_API_KEY)
with get_openai_callback() as cb:
result = llm.invoke("Tell me a joke")
print(cb)

Output:

'''
Tokens Used: 26
Prompt Tokens: 11
Completion Tokens: 15
Successful Requests: 1
Total Cost (USD): $0.00056
'''

Example - Tracking Multiple Calls:

with get_openai_callback() as cb:
result = llm.invoke("Tell me a joke")
result2 = llm.invoke("Tell me a joke")
print(cb.total_tokens)

Output:

'''
52
'''

Example - Tracking Token Usage with Multiple Calls in an Agent:

prompt = ChatPromptTemplate.from_messages(
[
("system", "You're a helpful assistant"),
("human", "{input}"),
("placeholder", "{agent_scratchpad}"),
]
)
tools = load_tools(["wikipedia"])
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(
agent=agent, tools=tools, verbose=True, stream_runnable=False
)

with get_openai_callback() as cb:
response = agent_executor.invoke(
{"input": "What's a hummingbird's scientific name and what's the fastest bird species?"}
)
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}")

Output:

'''

Entering new AgentExecutor chain...
...
Finished chain.
Total Tokens: 1790
Prompt Tokens: 1610
Completion Tokens: 180
Total Cost (USD): $0.0215
'''

In this episode, we delved into LangChain’s advanced features, focusing on how structured outputs, caching, and custom chat models can streamline your AI workflows. We also explored how tracking log probabilities and token usage can help you make data-driven decisions when fine-tuning your models.

By understanding and applying these features, you can develop more robust, efficient, and cost-effective AI applications that leverage the full potential of LLMs. Stay tuned for our next episode, where we’ll continue to uncover more tools and techniques to master LangChain.

---

Data Science Explore the world of data science with Donato_Story

Dashboard Discover the power of data visualization with Donato_Story

Donato_Journey Join me on my journey (Thai version)

Course_Review Discover the training courses with Donato_Story (Thai version)

I really value your thoughts and feedback. Please share any comments or questions you have below, or reach out to me on:

Medium: medium.com/donato-story
Facebook: web.facebook.com/DonatoStory
Linkedin: linkedin.com/in/nattapong-thanngam