ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/browser-use-rpa/browser-use/browser_use/llm/ollama/chat.py

@@ -0,0 +1,97 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, TypeVar, overload
+
+import httpx
+from ollama import AsyncClient as OllamaAsyncClient
+from ollama import Options
+from pydantic import BaseModel
+
+from browser_use.llm.base import BaseChatModel
+from browser_use.llm.exceptions import ModelProviderError
+from browser_use.llm.messages import BaseMessage
+from browser_use.llm.ollama.serializer import OllamaMessageSerializer
+from browser_use.llm.views import ChatInvokeCompletion
+
+T = TypeVar('T', bound=BaseModel)
+
+
+@dataclass
+class ChatOllama(BaseChatModel):
+	"""
+	A wrapper around Ollama's chat model.
+	"""
+
+	model: str
+
+	# # Model params
+	# TODO (matic): Why is this commented out?
+	# temperature: float | None = None
+
+	# Client initialization parameters
+	host: str | None = None
+	timeout: float | httpx.Timeout | None = None
+	client_params: dict[str, Any] | None = None
+	ollama_options: Mapping[str, Any] | Options | None = None
+
+	# Static
+	@property
+	def provider(self) -> str:
+		return 'ollama'
+
+	def _get_client_params(self) -> dict[str, Any]:
+		"""Prepare client parameters dictionary."""
+		return {
+			'host': self.host,
+			'timeout': self.timeout,
+			'client_params': self.client_params,
+		}
+
+	def get_client(self) -> OllamaAsyncClient:
+		"""
+		Returns an OllamaAsyncClient client.
+		"""
+		return OllamaAsyncClient(host=self.host, timeout=self.timeout, **self.client_params or {})
+
+	@property
+	def name(self) -> str:
+		return self.model
+
+	@overload
+	async def ainvoke(self, messages: list[BaseMessage], output_format: None = None) -> ChatInvokeCompletion[str]: ...
+
+	@overload
+	async def ainvoke(self, messages: list[BaseMessage], output_format: type[T]) -> ChatInvokeCompletion[T]: ...
+
+	async def ainvoke(
+		self, messages: list[BaseMessage], output_format: type[T] | None = None
+	) -> ChatInvokeCompletion[T] | ChatInvokeCompletion[str]:
+		ollama_messages = OllamaMessageSerializer.serialize_messages(messages)
+
+		try:
+			if output_format is None:
+				response = await self.get_client().chat(
+					model=self.model,
+					messages=ollama_messages,
+					options=self.ollama_options,
+				)
+
+				return ChatInvokeCompletion(completion=response.message.content or '', usage=None)
+			else:
+				schema = output_format.model_json_schema()
+
+				response = await self.get_client().chat(
+					model=self.model,
+					messages=ollama_messages,
+					format=schema,
+					options=self.ollama_options,
+				)
+
+				completion = response.message.content or ''
+				if output_format is not None:
+					completion = output_format.model_validate_json(completion)
+
+				return ChatInvokeCompletion(completion=completion, usage=None)
+
+		except Exception as e:
+			raise ModelProviderError(message=str(e), model=self.name) from e

chapter9/browser-use-rpa/browser-use/browser_use/llm/ollama/serializer.py

@@ -0,0 +1,143 @@
+import base64
+import json
+from typing import Any, overload
+
+from ollama._types import Image, Message
+
+from browser_use.llm.messages import (
+	AssistantMessage,
+	BaseMessage,
+	SystemMessage,
+	ToolCall,
+	UserMessage,
+)
+
+
+class OllamaMessageSerializer:
+	"""Serializer for converting between custom message types and Ollama message types."""
+
+	@staticmethod
+	def _extract_text_content(content: Any) -> str:
+		"""Extract text content from message content, ignoring images."""
+		if content is None:
+			return ''
+		if isinstance(content, str):
+			return content
+
+		text_parts: list[str] = []
+		for part in content:
+			if hasattr(part, 'type'):
+				if part.type == 'text':
+					text_parts.append(part.text)
+				elif part.type == 'refusal':
+					text_parts.append(f'[Refusal] {part.refusal}')
+			# Skip image parts as they're handled separately
+
+		return '\n'.join(text_parts)
+
+	@staticmethod
+	def _extract_images(content: Any) -> list[Image]:
+		"""Extract images from message content."""
+		if content is None or isinstance(content, str):
+			return []
+
+		images: list[Image] = []
+		for part in content:
+			if hasattr(part, 'type') and part.type == 'image_url':
+				url = part.image_url.url
+				if url.startswith('data:'):
+					# Handle base64 encoded images
+					# Format: data:image/png;base64,<data>
+					_, data = url.split(',', 1)
+					# Decode base64 to bytes
+					image_bytes = base64.b64decode(data)
+					images.append(Image(value=image_bytes))
+				else:
+					# Handle URL images (Ollama will download them)
+					images.append(Image(value=url))
+
+		return images
+
+	@staticmethod
+	def _serialize_tool_calls(tool_calls: list[ToolCall]) -> list[Message.ToolCall]:
+		"""Convert browser-use ToolCalls to Ollama ToolCalls."""
+		ollama_tool_calls: list[Message.ToolCall] = []
+
+		for tool_call in tool_calls:
+			# Parse arguments from JSON string to dict for Ollama
+			try:
+				arguments_dict = json.loads(tool_call.function.arguments)
+			except json.JSONDecodeError:
+				# If parsing fails, wrap in a dict
+				arguments_dict = {'arguments': tool_call.function.arguments}
+
+			ollama_tool_call = Message.ToolCall(
+				function=Message.ToolCall.Function(name=tool_call.function.name, arguments=arguments_dict)
+			)
+			ollama_tool_calls.append(ollama_tool_call)
+
+		return ollama_tool_calls
+
+	# region - Serialize overloads
+	@overload
+	@staticmethod
+	def serialize(message: UserMessage) -> Message: ...
+
+	@overload
+	@staticmethod
+	def serialize(message: SystemMessage) -> Message: ...
+
+	@overload
+	@staticmethod
+	def serialize(message: AssistantMessage) -> Message: ...
+
+	@staticmethod
+	def serialize(message: BaseMessage) -> Message:
+		"""Serialize a custom message to an Ollama Message."""
+
+		if isinstance(message, UserMessage):
+			text_content = OllamaMessageSerializer._extract_text_content(message.content)
+			images = OllamaMessageSerializer._extract_images(message.content)
+
+			ollama_message = Message(
+				role='user',
+				content=text_content if text_content else None,
+			)
+
+			if images:
+				ollama_message.images = images
+
+			return ollama_message
+
+		elif isinstance(message, SystemMessage):
+			text_content = OllamaMessageSerializer._extract_text_content(message.content)
+
+			return Message(
+				role='system',
+				content=text_content if text_content else None,
+			)
+
+		elif isinstance(message, AssistantMessage):
+			# Handle content
+			text_content = None
+			if message.content is not None:
+				text_content = OllamaMessageSerializer._extract_text_content(message.content)
+
+			ollama_message = Message(
+				role='assistant',
+				content=text_content if text_content else None,
+			)
+
+			# Handle tool calls
+			if message.tool_calls:
+				ollama_message.tool_calls = OllamaMessageSerializer._serialize_tool_calls(message.tool_calls)
+
+			return ollama_message
+
+		else:
+			raise ValueError(f'Unknown message type: {type(message)}')
+
+	@staticmethod
+	def serialize_messages(messages: list[BaseMessage]) -> list[Message]:
+		"""Serialize a list of browser_use messages to Ollama Messages."""
+		return [OllamaMessageSerializer.serialize(m) for m in messages]