Spaces:
Running
Running
File size: 3,590 Bytes
3cad23b c07f594 3cad23b c07f594 3cad23b c07f594 3cad23b c07f594 3cad23b c07f594 3cad23b c07f594 3cad23b c07f594 3cad23b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import json
# Import relevant functionality
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langchain_anthropic import ChatAnthropic
import sys
sys.path.append('.')
from toolformers.base import Tool as AgoraTool
from langchain_core.tools import tool as function_to_tool
from toolformers.base import StringParameter, Toolformer, Conversation
from utils import register_cost
COSTS = {
'claude-3-5-sonnet-latest': {
'input_tokens': 3e-6,
'output_tokens': 15e-6
},
'claude-3-5-haiku-latest': {
'input_tokens': 1e-6,
'output_tokens': 5e-6
}
}
class LangChainConversation(Conversation):
def __init__(self, model_name, agent, messages, category=None):
self.model_name = model_name
self.agent = agent
self.messages = messages
self.category = category
def chat(self, message, role='user', print_output=True) -> str:
self.messages.append(HumanMessage(content=message))
final_message = ''
aggregate = None
for chunk in self.agent.stream({"messages": self.messages}, stream_mode="values"):
print(chunk)
print("----")
for message in chunk['messages']:
if isinstance(message, AIMessage):
content = message.content
if isinstance(content, str):
final_message += content
else:
for content_chunk in content:
if isinstance(content_chunk, str):
final_message += content_chunk
aggregate = chunk if aggregate is None else (aggregate + chunk)
#final_message += chunk['agent']['messages'].content
total_cost = 0
for message in aggregate['messages']:
if isinstance(message, AIMessage):
for cost_name in ['input_tokens', 'output_tokens']:
total_cost += COSTS[self.model_name][cost_name] * message.usage_metadata[cost_name]
register_cost(self.category, total_cost)
self.messages.append(AIMessage(content=final_message))
#print(final_message)
return final_message
class LangChainAnthropicToolformer(Toolformer):
def __init__(self, model_name, api_key):
self.model_name = model_name
self.api_key = api_key
def new_conversation(self, prompt, tools, category=None):
tools = [function_to_tool(tool.as_annotated_function()) for tool in tools]
model = ChatAnthropic(model_name=self.model_name, api_key=self.api_key)
agent_executor = create_react_agent(model, tools)
return LangChainConversation(self.model_name, agent_executor, [SystemMessage(prompt)], category)
#weather_tool = AgoraTool("WeatherForecastAPI", "A simple tool that returns the weather", [StringParameter(
# name="location",
# description="The name of the location for which the weather forecast is requested.",
# required=True
#)], lambda location: 'Sunny', {
# "type": "string"
#})
#
#tools = [agora_tool_to_langchain(weather_tool)]
#toolformer = LangChainToolformer("claude-3-sonnet-20240229", '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')
#conversation = toolformer.new_conversation('You are a weather bot', [weather_tool])
#
#print(conversation.chat('What is the weather in San Francisco?'))
|