"""FabellaVLLM - LangChain BaseChatModel wrapping vLLM endpoint. Uses vLLM's native tool-calling pipeline for Gemma 4. The server is started with ``--enable-auto-tool-choice --tool-call-parser gemma4`` (see ``modal_app.py``), which makes vLLM parse the model's native ``<|tool_call>...`` markers into OpenAI-spec ``tool_calls`` JSON. This client passes the tool specs in OpenAI format and reads the parsed ``tool_calls`` straight off the response. """ import os import sys from typing import Any sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from langchain_core.language_models import BaseChatModel from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.outputs import ChatGeneration, ChatResult from pydantic import Field, PrivateAttr from openai import OpenAI class FabellaVLLM(BaseChatModel): """LangChain chat model backed by vLLM OpenAI-compatible API.""" base_url: str = Field(default="https://khoitruong071510--fabella-serve-drafter.modal.run") model_name: str = "gemma-4" temperature: float = 0.9 top_p: float = 0.95 max_tokens: int = 4096 seed: int = 0 _client: Any = PrivateAttr(default=None) _tools: list[dict] = PrivateAttr(default_factory=list) _tool_call_id: int = PrivateAttr(default=0) @property def _llm_type(self) -> str: return "fabella-vllm" @property def _identifying_params(self) -> dict: return { "base_url": self.base_url, "model_name": self.model_name, "temperature": self.temperature, "top_p": self.top_p, "max_tokens": self.max_tokens, "seed": self.seed, } def _get_client(self) -> OpenAI: if self._client is None: self._client = OpenAI( base_url=f"{self.base_url}/v1", api_key="EMPTY", ) return self._client def bind_tools(self, tools: list, **kwargs): # type: ignore[override] specs = [] for t in tools: specs.append(_to_openai_tool_spec(t)) object.__setattr__(self, "_tools", specs) object.__setattr__(self, "_tool_call_id", 0) return self def _generate(self, messages, stop=None, run_manager=None, **kwargs): client = self._get_client() system, non_system = _split_system(messages) api_messages = [] if system: api_messages.append({"role": "system", "content": system}) api_messages.extend(_to_api_messages(non_system)) request: dict[str, Any] = { "model": self.model_name, "messages": api_messages, "temperature": self.temperature, "top_p": self.top_p, "max_tokens": self.max_tokens, } if self.seed: request["seed"] = self.seed if self._tools: request["tools"] = self._tools response = client.chat.completions.create(**request) message = response.choices[0].message ai_message = _parse_response_message(message, state=self) return ChatResult(generations=[ChatGeneration(message=ai_message)]) def _split_system(messages) -> tuple[str, list]: system_parts: list[str] = [] rest: list = [] for m in messages: if isinstance(m, SystemMessage): content = m.content if isinstance(m.content, str) else str(m.content) system_parts.append(content) else: rest.append(m) return "\n".join(system_parts), rest def _to_api_messages(messages) -> list[dict]: """Translate LangChain messages to OpenAI chat-completions format.""" out: list[dict] = [] for m in messages: if isinstance(m, HumanMessage): content = m.content if isinstance(m.content, str) else str(m.content) out.append({"role": "user", "content": content}) elif isinstance(m, AIMessage): entry: dict[str, Any] = {"role": "assistant"} content = m.content if isinstance(m.content, str) else str(m.content) if content: entry["content"] = content if m.tool_calls: entry["tool_calls"] = [ { "id": tc.get("id", f"call_{i}"), "type": "function", "function": { "name": tc.get("name", ""), "arguments": _dump_args(tc.get("args", {})), }, } for i, tc in enumerate(m.tool_calls) ] out.append(entry) elif isinstance(m, ToolMessage): content = m.content if isinstance(m.content, str) else str(m.content) entry = { "role": "tool", "tool_call_id": m.tool_call_id, "content": content, } out.append(entry) else: content = getattr(m, "content", "") content = content if isinstance(content, str) else str(content) out.append({"role": "user", "content": content}) return out def _parse_response_message(message, *, state: "FabellaVLLM") -> AIMessage: content = message.content or "" if not message.tool_calls: return AIMessage(content=content) tool_calls = [] for tc in message.tool_calls: state._tool_call_id += 1 raw_args = tc.function.arguments args = _loads_args(raw_args) tool_calls.append( { "name": tc.function.name, "args": args, "id": tc.id or f"call_{state._tool_call_id}", "type": "tool_call", } ) return AIMessage(content=content, tool_calls=tool_calls) def _to_openai_tool_spec(tool_obj) -> dict: """Build an OpenAI-spec tool entry from a LangChain tool.""" name = getattr(tool_obj, "name", None) or getattr(tool_obj, "__name__", "tool") description = (getattr(tool_obj, "description", "") or (tool_obj.__doc__ or "")).strip() parameters = _extract_parameters(tool_obj) return { "type": "function", "function": { "name": name, "description": description, "parameters": parameters, }, } def _extract_parameters(tool_obj) -> dict: try: from langchain_core.tools import BaseTool if isinstance(tool_obj, BaseTool): schema = tool_obj.args properties = { name: _normalize_schema(field) for name, field in schema.items() } required = [ name for name, field in schema.items() if field.get("type") != "null" and name not in (schema.get("additionalProperties") or {}) ] return { "type": "object", "properties": properties, "required": list(schema.keys()), } except Exception: pass if hasattr(tool_obj, "args_schema") and tool_obj.args_schema is not None: try: model = tool_obj.args_schema from pydantic import BaseModel # type: ignore if isinstance(model, type) and issubclass(model, BaseModel): return model.model_json_schema() if hasattr(model, "model_json_schema"): return model.model_json_schema() if hasattr(model, "schema"): return model.schema() except Exception: pass return {"type": "object", "properties": {}} def _normalize_schema(field: dict) -> dict: out = {"type": field.get("type", "string")} if "description" in field: out["description"] = field["description"] if "enum" in field: out["enum"] = field["enum"] return out def _dump_args(args: Any) -> str: import json if isinstance(args, str): return args return json.dumps(args, ensure_ascii=False) def _loads_args(raw: Any) -> Any: import json if isinstance(raw, dict): return raw if not raw: return {} try: return json.loads(raw) except (TypeError, ValueError): return {"input": raw}