Skip to content

Commit

Permalink
Print prompt caching data in example
Browse files Browse the repository at this point in the history
  • Loading branch information
vblagoje committed Aug 20, 2024
1 parent 2db1616 commit 37908ce
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
4 changes: 3 additions & 1 deletion integrations/anthropic/example/prompt_caching.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,11 +53,13 @@
# Answer the questions using prompt caching (i.e. the entire document is cached, we run the question against it)
for question in questions:
print("Question: " + question)
qa_pipeline.run(
result = qa_pipeline.run(
data={
"llm": {"messages": [final_prompt_msg,
ChatMessage.from_user("Given these documents, answer the question:" + question)]},
}
)

print("\n\nChecking cache usage:", result["llm"]["replies"][0].meta.get("usage"))
print("\n")

Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ def run(self, messages: List[ChatMessage], generation_kwargs: Optional[Dict[str,
anthropic_formatted_messages = self._convert_to_anthropic_format(messages)

# system message provided by the user overrides the system message from the self.generation_kwargs
system = messages[0].content if messages and messages[0].is_from(ChatRole.SYSTEM) else None
system = [anthropic_formatted_messages[0]] if messages and messages[0].is_from(ChatRole.SYSTEM) else None
if system:
anthropic_formatted_messages = anthropic_formatted_messages[1:]

Expand Down Expand Up @@ -262,6 +262,11 @@ def _convert_to_anthropic_format(self, messages: List[ChatMessage]) -> List[Dict
for m in messages:
message_dict = dataclasses.asdict(m)
filtered_message = {k: v for k, v in message_dict.items() if k in {"role", "content"} and v}
if m.is_from(ChatRole.SYSTEM):
# system messages need to be in the format expected by the Anthropic API
filtered_message.pop("role")
filtered_message["type"] = "text"
filtered_message["text"] = filtered_message.pop("content")
filtered_message.update(m.meta or {})
anthropic_formatted_messages.append(filtered_message)
return anthropic_formatted_messages
Expand Down

0 comments on commit 37908ce

Please sign in to comment.