diff --git a/src/engram/__init__.py b/src/engram/__init__.py index 5e6eee2..9fbb4b5 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -1,15 +1,16 @@ from ._models import ( CommittedOperation, CommittedOperations, - ConversationContent, + ConversationInput, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, Run, RunStatus, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -33,18 +34,19 @@ "CommittedOperation", "CommittedOperations", "ConnectionError", - "ConversationContent", + "ConversationInput", "EngramClient", "EngramError", "EngramTimeoutError", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index 06656aa..b9906e2 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -1,12 +1,13 @@ from .memory import ( - AddContent, - ConversationContent, + AddInput, + ConversationInput, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -14,18 +15,19 @@ from .run import CommittedOperation, CommittedOperations, Run, RunStatus __all__ = [ - "AddContent", + "AddInput", "CommittedOperation", "CommittedOperations", - "ConversationContent", + "ConversationInput", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 6b21f58..17ed194 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -6,18 +6,27 @@ @dataclass(slots=True) -class PreExtractedContent: - """Pre-extracted content that bypasses the extraction pipeline.""" +class PreExtractedInput: + """Pre-extracted input that skips the extraction step continues through the pipeline as-is. + Each individual item represents a separate memory. + """ + + items: list[PreExtractedItem] + + +@dataclass(slots=True) +class PreExtractedItem: + """A single pre-extracted memory.""" content: str topic: str @dataclass(slots=True) -class StringContent: - """String content that bypasses the extraction pipeline.""" +class StringInput: + """String input to extract memories from.""" - content: str + content: str | list[str] @dataclass(slots=True) @@ -50,7 +59,7 @@ class ToolCallInput: @dataclass(slots=True) -class MessageContent: +class MessageInput: """A message in a conversation using the OpenAI Chat Completions format. - 'tool' role (tool results) is mapped to 'user' by the server. @@ -66,18 +75,18 @@ class MessageContent: @dataclass(slots=True) -class ConversationContent: - """Conversation content that bypasses the extraction pipeline.""" +class ConversationInput: + """Conversation input that bypasses the extraction pipeline.""" - messages: list[MessageContent] + messages: list[MessageInput] metadata: dict[str, Any] | None = None created_at: str | None = None updated_at: str | None = None -# Type alias for the content argument to memories.add() -AddContent: TypeAlias = ( - str | list[dict[str, str]] | PreExtractedContent | ConversationContent | StringContent +# Type alias for the input_data argument to memories.add() +AddInput: TypeAlias = ( + str | list[dict[str, str]] | PreExtractedInput | ConversationInput | StringInput ) diff --git a/src/engram/_resources/memories.py b/src/engram/_resources/memories.py index fa0bcba..1b22627 100644 --- a/src/engram/_resources/memories.py +++ b/src/engram/_resources/memories.py @@ -3,7 +3,7 @@ from uuid import UUID from .._http import AsyncHttpTransport, HttpTransport -from .._models import AddContent, Memory, RetrievalConfig, Run, SearchResults +from .._models import AddInput, Memory, RetrievalConfig, Run, SearchResults from .._serialization import ( build_add_body, build_memory_params, @@ -29,14 +29,14 @@ def __init__(self, transport: HttpTransport) -> None: def add( self, - content: AddContent, + input_data: AddInput, *, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, ) -> Run: body = build_add_body( - content, + input_data, user_id=user_id, conversation_id=conversation_id, group=group, @@ -101,14 +101,14 @@ def __init__(self, transport: AsyncHttpTransport) -> None: async def add( self, - content: AddContent, + input_data: AddInput, *, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, ) -> Run: body = build_add_body( - content, + input_data, user_id=user_id, conversation_id=conversation_id, group=group, diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index abcd0c5..7a975db 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -3,11 +3,11 @@ from typing import Any from .._models import ( - AddContent, - ConversationContent, - PreExtractedContent, + AddInput, + ConversationInput, + PreExtractedInput, RetrievalConfig, - StringContent, + StringInput, ToolCallInput, ) @@ -21,29 +21,28 @@ def _serialize_tool_call(tc: ToolCallInput) -> dict[str, Any]: return out -def _serialize_content(content: AddContent) -> dict[str, Any]: - """Build the content envelope with the type discriminator.""" - if isinstance(content, str): - return {"type": "string", "content": content} - if isinstance(content, StringContent): - return {"type": "string", "content": content.content} - if isinstance(content, PreExtractedContent): +def _serialize_input(input_data: AddInput) -> dict[str, Any]: + """Build the input envelope with the type discriminator.""" + if isinstance(input_data, str): + return {"string": {"content": [input_data]}} + if isinstance(input_data, StringInput): + if isinstance(input_data.content, list): + return {"string": {"content": input_data.content}} + else: + return {"string": {"content": [input_data.content]}} + if isinstance(input_data, PreExtractedInput): + items = [{"content": item.content, "topic": item.topic} for item in input_data.items] + return {"pre_extracted": {"items": items}} + if isinstance(input_data, list): return { - "type": "pre_extracted", - "content": content.content, - "topic": content.topic, + "conversation": {"messages": input_data}, } - if isinstance(content, list): - return { - "type": "conversation", - "conversation": {"messages": content}, - } - if isinstance(content, ConversationContent): - return _serialize_conversation_content(content) - raise TypeError(f"Unsupported content type: {type(content)}") # pragma: no cover + if isinstance(input_data, ConversationInput): + return _serialize_conversation_content(input_data) + raise TypeError(f"Unsupported input type: {type(input_data)}") # pragma: no cover -def _serialize_conversation_content(content: ConversationContent) -> dict[str, Any]: +def _serialize_conversation_content(content: ConversationInput) -> dict[str, Any]: messages = [] for msg in content.messages: m: dict[str, Any] = {"role": msg.role, "content": msg.content} @@ -63,17 +62,17 @@ def _serialize_conversation_content(content: ConversationContent) -> dict[str, A conversation["created_at"] = content.created_at if content.updated_at is not None: conversation["updated_at"] = content.updated_at - return {"type": "conversation", "conversation": conversation} + return {"conversation": conversation} def build_add_body( - content: AddContent, + input_data: AddInput, *, user_id: str | None, conversation_id: str | None, group: str | None, ) -> dict[str, Any]: - body: dict[str, Any] = {"content": _serialize_content(content)} + body: dict[str, Any] = {"input": _serialize_input(input_data)} if user_id is not None: body["user_id"] = user_id if conversation_id is not None: diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 6942072..64d35f8 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -6,11 +6,12 @@ from engram._http import AsyncHttpTransport from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallFuncInput, ToolCallInput, ) @@ -117,7 +118,7 @@ async def test_add_str() -> None: async def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = await client.memories.add( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -146,16 +147,65 @@ def handler(request: httpx.Request) -> httpx.Response: await client.memories.add("hello", user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } +@pytest.mark.asyncio +async def test_add_multiple_strings() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + await client.memories.add(StringInput(content=["first", "second"]), user_id="u1") + body = json.loads(captured[0].content) + assert body == { + "input": {"string": {"content": ["first", "second"]}}, + "user_id": "u1", + } + + +@pytest.mark.asyncio +async def test_add_multiple_pre_extracted_items() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + await client.memories.add( + PreExtractedInput( + items=[ + PreExtractedItem(content="fact one", topic="topic_a"), + PreExtractedItem(content="fact two", topic="topic_b"), + ] + ), + user_id="u1", + ) + body = json.loads(captured[0].content) + assert body == { + "input": { + "pre_extracted": { + "items": [ + {"content": "fact one", "topic": "topic_a"}, + {"content": "fact two", "topic": "topic_b"}, + ] + } + }, + "user_id": "u1", + } + + @pytest.mark.asyncio async def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) - result = await client.memories.add(StringContent(content="hello"), user_id="u1") + result = await client.memories.add(StringInput(content="hello"), user_id="u1") assert result.run_id == "r4" @@ -168,10 +218,10 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - await client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") + await client.memories.add(StringInput(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -181,7 +231,7 @@ def handler(request: httpx.Request) -> httpx.Response: async def test_add_conversation_content() -> None: client = _make_client(body={"run_id": "r5", "status": "pending"}) result = await client.memories.add( - ConversationContent(messages=[MessageContent(role="user", content="hi")]), + ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", conversation_id="c1", ) @@ -198,10 +248,10 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) await client.memories.add( - ConversationContent( + ConversationInput( messages=[ - MessageContent(role="user", content="hi"), - MessageContent( + MessageInput(role="user", content="hi"), + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -215,8 +265,7 @@ def handler(request: httpx.Request) -> httpx.Response: conversation_id="c1", ) body = json.loads(captured[0].content) - assert body["content"]["type"] == "conversation" - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index fe6586b..5e666c9 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -6,11 +6,12 @@ from engram._http import HttpTransport from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallFuncInput, ToolCallInput, ) @@ -118,7 +119,7 @@ def test_add_str() -> None: def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = client.memories.add( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -145,7 +146,7 @@ def handler(request: httpx.Request) -> httpx.Response: client.memories.add("hello", user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -163,17 +164,61 @@ def handler(request: httpx.Request) -> httpx.Response: client.memories.add(messages, conversation_id="c1") body = json.loads(captured[0].content) assert body == { - "content": { - "type": "conversation", - "conversation": {"messages": messages}, - }, + "input": {"conversation": {"messages": messages}}, "conversation_id": "c1", } +def test_add_multiple_strings() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + client.memories.add(StringInput(content=["first", "second"]), user_id="u1") + body = json.loads(captured[0].content) + assert body == { + "input": {"string": {"content": ["first", "second"]}}, + "user_id": "u1", + } + + +def test_add_multiple_pre_extracted_items() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + client.memories.add( + PreExtractedInput( + items=[ + PreExtractedItem(content="fact one", topic="topic_a"), + PreExtractedItem(content="fact two", topic="topic_b"), + ] + ), + user_id="u1", + ) + body = json.loads(captured[0].content) + assert body == { + "input": { + "pre_extracted": { + "items": [ + {"content": "fact one", "topic": "topic_a"}, + {"content": "fact two", "topic": "topic_b"}, + ] + } + }, + "user_id": "u1", + } + + def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) - result = client.memories.add(StringContent(content="hello"), user_id="u1") + result = client.memories.add(StringInput(content="hello"), user_id="u1") assert result.run_id == "r4" @@ -185,10 +230,10 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") + client.memories.add(StringInput(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -197,7 +242,7 @@ def handler(request: httpx.Request) -> httpx.Response: def test_add_conversation_content() -> None: client = _make_client(body={"run_id": "r5", "status": "pending"}) result = client.memories.add( - ConversationContent(messages=[MessageContent(role="user", content="hi")]), + ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", conversation_id="c1", ) @@ -213,10 +258,10 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) client.memories.add( - ConversationContent( + ConversationInput( messages=[ - MessageContent(role="user", content="hi"), - MessageContent( + MessageInput(role="user", content="hi"), + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -230,8 +275,7 @@ def handler(request: httpx.Request) -> httpx.Response: conversation_id="c1", ) body = json.loads(captured[0].content) - assert body["content"]["type"] == "conversation" - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} diff --git a/tests/test_imports.py b/tests/test_imports.py index 0a8e8e4..fb29e1f 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -7,18 +7,19 @@ def test_public_imports() -> None: CommittedOperation, CommittedOperations, ConnectionError, - ConversationContent, + ConversationInput, EngramClient, EngramError, EngramTimeoutError, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, Run, RunStatus, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -36,13 +37,14 @@ def test_public_imports() -> None: assert isinstance(Run, type) assert isinstance(RunStatus, type) assert isinstance(SearchResults, type) - assert isinstance(PreExtractedContent, type) + assert isinstance(PreExtractedInput, type) + assert isinstance(PreExtractedItem, type) assert isinstance(RetrievalConfig, type) assert isinstance(CommittedOperation, type) assert isinstance(CommittedOperations, type) - assert isinstance(ConversationContent, type) - assert isinstance(MessageContent, type) - assert isinstance(StringContent, type) + assert isinstance(ConversationInput, type) + assert isinstance(MessageInput, type) + assert isinstance(StringInput, type) assert isinstance(ToolCallCustomInput, type) assert isinstance(ToolCallFuncInput, type) assert isinstance(ToolCallInput, type) @@ -54,18 +56,19 @@ def test_public_imports() -> None: "CommittedOperation", "CommittedOperations", "ConnectionError", - "ConversationContent", + "ConversationInput", "EngramClient", "EngramError", "EngramTimeoutError", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 3fe1ab2..6328d81 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -1,9 +1,10 @@ from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, + PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -28,7 +29,7 @@ def test_build_add_body_str() -> None: conversation_id=None, group=None, ) - assert body == {"content": {"type": "string", "content": "hello world"}} + assert body == {"input": {"string": {"content": ["hello world"]}}} def test_build_add_body_str_with_options() -> None: @@ -39,7 +40,7 @@ def test_build_add_body_str_with_options() -> None: group="g1", ) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "conversation_id": "c1", "group": "g1", @@ -48,13 +49,13 @@ def test_build_add_body_str_with_options() -> None: def test_build_add_body_pre_extracted() -> None: body = build_add_body( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id=None, conversation_id=None, group=None, ) assert body == { - "content": {"type": "pre_extracted", "content": "fact", "topic": "topic"}, + "input": {"pre_extracted": {"items": [{"content": "fact", "topic": "topic"}]}}, } @@ -70,10 +71,7 @@ def test_build_add_body_conversation() -> None: group=None, ) assert body == { - "content": { - "type": "conversation", - "conversation": {"messages": messages}, - }, + "input": {"conversation": {"messages": messages}}, "user_id": "u1", "conversation_id": "c1", } @@ -81,23 +79,23 @@ def test_build_add_body_conversation() -> None: def test_build_add_body_string_content() -> None: body = build_add_body( - StringContent(content="hello world"), + StringInput(content="hello world"), user_id=None, conversation_id=None, group=None, ) - assert body == {"content": {"type": "string", "content": "hello world"}} + assert body == {"input": {"string": {"content": ["hello world"]}}} def test_build_add_body_string_content_with_options() -> None: body = build_add_body( - StringContent(content="hello"), + StringInput(content="hello"), user_id="u1", conversation_id="c1", group="g1", ) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "conversation_id": "c1", "group": "g1", @@ -106,18 +104,17 @@ def test_build_add_body_string_content_with_options() -> None: def test_build_add_body_conversation_content() -> None: messages = [ - MessageContent(role="user", content="hi"), - MessageContent(role="assistant", content="hello"), + MessageInput(role="user", content="hi"), + MessageInput(role="assistant", content="hello"), ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id="u1", conversation_id="c1", group=None, ) assert body == { - "content": { - "type": "conversation", + "input": { "conversation": { "messages": [ {"role": "user", "content": "hi"}, @@ -131,9 +128,9 @@ def test_build_add_body_conversation_content() -> None: def test_build_add_body_conversation_content_with_metadata() -> None: - messages = [MessageContent(role="user", content="hi")] + messages = [MessageInput(role="user", content="hi")] body = build_add_body( - ConversationContent( + ConversationInput( messages=messages, metadata={"session_id": "s1"}, created_at="2024-01-01T00:00:00Z", @@ -143,28 +140,28 @@ def test_build_add_body_conversation_content_with_metadata() -> None: conversation_id=None, group=None, ) - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["created_at"] == "2024-01-01T00:00:00Z" assert conv["updated_at"] == "2024-01-02T00:00:00Z" def test_build_add_body_conversation_content_with_message_timestamps() -> None: - messages = [MessageContent(role="user", content="hi", created_at="2024-01-01T00:00:00Z")] + messages = [MessageInput(role="user", content="hi", created_at="2024-01-01T00:00:00Z")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["created_at"] == "2024-01-01T00:00:00Z" assert "tool_call_metadata" not in msg def test_build_add_body_conversation_content_with_tool_calls() -> None: messages = [ - MessageContent( + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -174,12 +171,12 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: ) ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": '{"q":"x"}'}} ] @@ -187,7 +184,7 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: messages = [ - MessageContent( + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -199,26 +196,26 @@ def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: ) ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["tool_calls"] == [ {"id": "tc2", "type": "custom", "custom": {"name": "my_tool", "input": "some input"}} ] def test_build_add_body_conversation_content_with_tool_role() -> None: - messages = [MessageContent(role="tool", content="result", tool_call_id="tc1", name="search")] + messages = [MessageInput(role="tool", content="result", tool_call_id="tc1", name="search")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["role"] == "tool" assert msg["tool_call_id"] == "tc1" assert msg["name"] == "search" @@ -226,14 +223,14 @@ def test_build_add_body_conversation_content_with_tool_role() -> None: def test_build_add_body_conversation_content_with_developer_role() -> None: - messages = [MessageContent(role="developer", content="You are a helpful assistant.")] + messages = [MessageInput(role="developer", content="You are a helpful assistant.")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["role"] == "developer" assert msg["content"] == "You are a helpful assistant."