Metadata-Version: 2.4
Name: custom-langchain-model
Version: 0.2.0
Summary: LangChain custom chat model
Author: Thanh Huynh
License: MIT
Project-URL: Homepage, https://github.com/thanhhuynhk17/ChatBaml
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: baml-py==0.220.0
Requires-Dist: langchain>=1.0.0
Requires-Dist: langgraph>=1.0.0
Requires-Dist: pydantic>=2.12.3
Requires-Dist: pydantic-settings>=2.11.0
Dynamic: license-file

# ChatBaml

Custom LangChain Chat Model tích hợp BAML — bridge giữa LangChain orchestration và BAML structured extraction/function calling.

> **Demo**: React Agent với ChatBaml chạy trên Qwen3-VL-8B-Instruct
> ![demo](images/chatbaml_react_agent.gif)

---

## Mục lục

- [Cài đặt](#cài-đặt)
- [Cấu hình môi trường](#cấu-hình-môi-trường)
- [Kiến trúc](#kiến-trúc)
- [Sử dụng](#sử-dụng)
  - [1. Standalone BAML Client](#1-standalone-baml-client)
  - [2. Pydantic Extraction](#2-pydantic-extraction)
  - [3. ChatBaml với LangChain](#3-chatbaml-với-langchain)
  - [4. React Agent](#4-react-agent)
  - [5. Tool Definition](#5-tool-definition)
- [BAML Functions](#baml-functions)
- [API Reference](#api-reference)

---

## Cài đặt

```bash
git clone https://github.com/thanhhuynhk17/ChatBaml.git
cd ChatBaml

uv venv
source .venv/bin/activate          # Linux/macOS
# source .venv/Scripts/activate    # Windows

uv pip install -e .

# Generate baml_client/ từ baml_src/
uv run baml-cli generate
```

---

## Cấu hình môi trường

Tạo file `.env` (xem `.env.sample`):

```bash
# Model config
OPENAI_MODEL_NAME="qwen3-vl"
OPENAI_BASE_URL="http://localhost:8000/v1"
OPENAI_API_KEY="sk-your-key"

# Bắt buộc khi dùng vLLM
DEFAULT_ROLE="user"

# Tắt BAML verbose logging (optional)
BAML_LOG=off

# Bắt buộc cho LangChain content blocks v1
LC_OUTPUT_VERSION="v1"
```

---

## Kiến trúc

```
custom_langchain_model/
├── core/
│   ├── registry.py        ← build_client_registry(), get_baml_client()
│   └── logging.py
├── llms/
│   ├── chat_baml.py       ← ChatBaml (BaseChatModel)
│   └── types.py
├── helpers/
│   ├── parse_json_schema.py  ← convert_to_baml_tool()
│   ├── messages.py
│   └── render_agent_wants_to.py
└── extractors/
    └── structured.py      ← extract_structured(), extract_with_pydantic()

baml_src/
├── chat_baml.baml         ← Chat, ChooseTool, ExtractStructure
├── clients.baml
└── generators.baml

react_agent/
├── agent.py
└── tools.py
```

**Nguyên tắc thiết kế:**
- `core/registry.py` là **single source of truth** cho BAML authentication — cả `ChatBaml` lẫn standalone client đều dùng chung.
- Không cần tạo `ChatBaml` instance nếu chỉ cần gọi BAML functions trực tiếp.

---

## Sử dụng

### 1. Standalone BAML Client

Dùng khi muốn gọi BAML functions trực tiếp mà **không cần** tạo LangChain model instance. Giải quyết vấn đề 401 auth error khi dùng `baml_client.b` mặc định.

```python
from custom_langchain_model import get_baml_client

# Đọc config từ env vars (OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL_NAME)
b = get_baml_client()

# Override bất kỳ param nào
b = get_baml_client(
    model="gpt-4o",
    api_key="sk-...",
    base_url="http://localhost:8000/v1",
)

# Gọi BAML functions như bình thường
from custom_langchain_model.baml_client.types import BamlState, BaseMessage, ContentBlock

state = BamlState(messages=[
    BaseMessage(role="user", content_block=ContentBlock(text="Hello!"))
])
result = b.Chat(state)
print(result)  # str
```

---

### 2. Pydantic Extraction — `extract_with_pydantic()`

Trích xuất dữ liệu sử dụng Pydantic Model. Hỗ trợ nested models và truyền `description` trong `Field` để điều hướng LLM.

```python
from pydantic import BaseModel, Field
from custom_langchain_model import extract_with_pydantic

class Invoice(BaseModel):
    invoice_number: str = Field(description="Số hóa đơn")
    total_amount: float = Field(description="Tổng tiền")
    customer_name: str = Field(description="Tên khách hàng, hãy viết HOA")

result = extract_with_pydantic(
    input_text="Hóa đơn #INV-2024-001 cho khách Nguyễn Văn A, tổng 5.500.000đ",
    model_class=Invoice,
    client_kwargs={
        "base_url": "http://localhost:8000/v1",
        "api_key": "sk-...",
        "model": "qwen3-vl"
    }
)
print(result.extracted["customer_name"])  # "NGUYỄN VĂN A"
```

**Lưu ý**: Sử dụng `client_kwargs` để override cấu hình client (base\_url, api\_key, model) thay vì dùng file `.env`.


---

### 3. ChatBaml với LangChain

```python
import os
from dotenv import load_dotenv
from custom_langchain_model import ChatBaml

load_dotenv()

chat = ChatBaml(
    base_url=os.getenv("OPENAI_BASE_URL"),
    api_key=os.getenv("OPENAI_API_KEY"),
    model=os.getenv("OPENAI_MODEL_NAME"),
    temperature=0.7,
    streaming=True,
)

# Basic invoke
from langchain_core.messages import HumanMessage, SystemMessage
response = chat.invoke([
    SystemMessage(content="Bạn là trợ lý hữu ích."),
    HumanMessage(content="2 + 2 bằng mấy?"),
])
print(response.content)

# Streaming
for chunk in chat.stream([HumanMessage(content="Kể tên 3 thủ đô châu Á.")]):
    print(chunk.content, end="", flush=True)
```

#### Bind tools

```python
from pydantic import BaseModel, Field
from custom_langchain_model import ChatBaml

class SearchWeb(BaseModel):
    """Tìm kiếm thông tin trên internet."""
    query: str = Field(description="Từ khóa tìm kiếm")

class Calculator(BaseModel):
    """Tính toán biểu thức số học."""
    expression: str = Field(description="Biểu thức cần tính, ví dụ: '2 + 3 * 4'")

chat = ChatBaml(model="qwen3-vl", streaming=True)
chat_with_tools = chat.bind_tools([SearchWeb, Calculator])

response = chat_with_tools.invoke([
    HumanMessage(content="Tìm kiếm về LangChain trên internet")
])

if response.tool_calls:
    tool_call = response.tool_calls[0]
    print(f"Tool: {tool_call['name']}")
    print(f"Args: {tool_call['args']}")
```

---

### 4. React Agent

```bash
# Chạy LangGraph dev server
langgraph dev

# Hoặc F5 trong VSCode (cần .vscode/launch.json)
```

```python
# react_agent/agent.py
from langchain.agents import create_agent
from custom_langchain_model import ChatBaml
from react_agent.tools import add, multiply
from langchain_community.tools import DuckDuckGoSearchRun

chat_baml = ChatBaml(
    base_url=os.getenv("OPENAI_BASE_URL"),
    api_key=os.getenv("OPENAI_API_KEY"),
    model=os.getenv("OPENAI_MODEL_NAME"),
    streaming=True,
    temperature=0.7,
)

agent = create_agent(
    chat_baml,
    tools=[add, multiply, DuckDuckGoSearchRun()],
    system_prompt="You are a helpful assistant with math and search tools.",
)
```

---

### 5. Tool Definition

#### Pydantic BaseModel (recommended)

```python
from pydantic import BaseModel, Field

class AddTool(BaseModel):
    """Cộng hai số nguyên."""
    a: int = Field(..., description="Số thứ nhất")
    b: int = Field(..., description="Số thứ hai")

class SearchTool(BaseModel):
    """Tìm kiếm thông tin."""
    query: str = Field(..., description="Từ khóa tìm kiếm")
    max_results: int = Field(default=5, description="Số kết quả tối đa")
```

#### LangChain `@tool` (cần `parse_docstring=True`)

```python
from langchain.tools import tool

@tool(parse_docstring=True)
def get_weather(city: str, unit: str = "celsius") -> str:
    """Lấy thông tin thời tiết hiện tại của một thành phố.

    Args:
        city (str): Tên thành phố.
        unit (str): Đơn vị nhiệt độ, "celsius" hoặc "fahrenheit".

    Returns:
        str: Thông tin thời tiết dạng text.
    """
    return f"Thời tiết tại {city}: 28°C, nhiều mây"
```

> **Lưu ý**: `@tool` phải có docstring đầy đủ `Args:` và `Returns:` sections. Thiếu sẽ gây lỗi khi convert sang BAML schema.

---

## BAML Functions

| Function          | Input                                       | Output                    | Dùng khi                                  |
|-------------------|---------------------------------------------|---------------------------|-------------------------------------------|
| `Chat`            | `BamlState`                                 | `string`                  | Chat thông thường, không có tools         |
| `ChooseTool`      | `BamlState` + `TypeBuilder`                 | `string \| DynamicSchema` | Agent cần chọn và gọi tools               |
| `ExtractStructure`| `system_instruction?`, `prefix?`, `input`   | `DynamicSchema`           | Trích xuất structured data từ text        |

### ExtractStructure chi tiết

```baml
function ExtractStructure(
    system_instruction: string?,   // null → dùng default instruction
    prefix: string?,               // null → không có prefix trong output format
    input: string                  // text cần extract (bắt buộc)
) -> DynamicSchema
```

**Cách BAML render prompt:**
- `system_instruction | default(...)` → fallback về default nếu null
- `{% if prefix %}` → chỉ thêm prefix vào `ctx.output_format` khi có giá trị
- `DynamicSchema` được define lúc runtime qua `TypeBuilder` (truyền vào `baml_options={"tb": tb}`)

---

## API Reference

### `get_baml_client(*, model, api_key, base_url, **kwargs)`

Trả về BAML client đã authenticated. Không cần tạo `ChatBaml`.

```python
from custom_langchain_model import get_baml_client
b = get_baml_client()
```

### `extract_with_pydantic(input_text, model_class, client_kwargs=None, ...)`

Extraction với Pydantic model làm schema. Hỗ trợ `client_kwargs` để cấu hình client.

```python
from custom_langchain_model import extract_with_pydantic
result = extract_with_pydantic("text...", MyModel, client_kwargs={"model": "gpt-4o"})
```

### `ChatBaml`

LangChain `BaseChatModel` backed by BAML. Hỗ trợ `.invoke()`, `.stream()`, `.bind_tools()`.

```python
from custom_langchain_model import ChatBaml
chat = ChatBaml(model="qwen3-vl", streaming=True)
```

---

## References

- [Original repo](https://github.com/tranngocphu/custom_langchain_chat_model)
- [BAML Docs](https://docs.boundaryml.com)
- [LangChain Custom Models](https://docs.langchain.com/langsmith/custom-endpoint)
- [LangGraph](https://langchain-ai.github.io/langgraph/)
