Initial upload to test HF Spaces.

.gitignore

@@ -0,0 +1,4 @@
+venv/**
+venv2/**
+*.pyc
+.env

app.py

@@ -0,0 +1,276 @@
+import json
+
+import gradio as gr
+
+from collections import UserList
+
+from flow import full_flow
+
+schema = {
+  "input": {
+    "type": "object",
+    "properties": {
+      "location": {
+        "type": "string",
+        "description": "The name of the location for which the weather forecast is requested."
+      },
+      "date": {
+        "type": "string",
+        "format": "date",
+        "description": "The date for which the weather forecast is requested, in YYYY-MM-DD format."
+      }
+    },
+    "required": [
+      "location",
+      "date"
+    ]
+  },
+  "output": {
+    "type": "object",
+    "properties": {
+      "temperature": {
+        "type": "number",
+        "description": "The forecasted temperature in degrees Celsius."
+      },
+      "condition": {
+        "type": "string",
+        "description": "A brief description of the weather condition (e.g., sunny, cloudy, rainy)."
+      },
+      "humidity": {
+        "type": "number",
+        "description": "The forecasted humidity percentage."
+      },
+      "wind_speed": {
+        "type": "number",
+        "description": "The forecasted wind speed in kilometers per hour."
+      }
+    },
+    "required": [
+      "temperature",
+      "condition",
+      "humidity",
+      "wind_speed"
+    ]
+  },
+  "description": "Alice requests a weather forecast for a specific location and date from Bob's weather service.",
+  "examples": [
+    {
+      "location": "New York",
+      "date": "2023-10-15"
+    },
+    {
+      "location": "London",
+      "date": "2023-11-01"
+    }
+  ],
+  "tools": [
+    {
+      "name": "WeatherForecastAPI",
+      "description": "An API that provides weather forecasts for a given location and date.",
+      "input": {
+        "type": "object",
+        "properties": {
+          "location": {
+            "type": "string",
+            "description": "The name of the location for which the weather forecast is requested."
+          },
+          "date": {
+            "type": "string",
+            "format": "date",
+            "description": "The date for which the weather forecast is requested, in YYYY-MM-DD format."
+          }
+        },
+        "required": [
+          "location",
+          "date"
+        ]
+      },
+      "output": {
+        "type": "object",
+        "properties": {
+          "temperature": {
+            "type": "number",
+            "description": "The forecasted temperature in degrees Celsius."
+          },
+          "condition": {
+            "type": "string",
+            "description": "A brief description of the weather condition (e.g., sunny, cloudy, rainy)."
+          },
+          "humidity": {
+            "type": "number",
+            "description": "The forecasted humidity percentage."
+          },
+          "wind_speed": {
+            "type": "number",
+            "description": "The forecasted wind speed in kilometers per hour."
+          }
+        },
+        "required": [
+          "temperature",
+          "condition",
+          "humidity",
+          "wind_speed"
+        ]
+      },
+      "dummy_outputs": [
+        {
+          "temperature": 18,
+          "condition": "Sunny",
+          "humidity": 55,
+          "wind_speed": 10
+        },
+        {
+          "temperature": 12,
+          "condition": "Cloudy",
+          "humidity": 80,
+          "wind_speed": 15
+        }
+      ]
+    }
+  ]
+}
+
+SCHEMAS = {
+    "weather_forecast": schema,
+    "other": { "input": "PIPPO"}
+}
+
+def parse_raw_messages(messages_raw):
+    messages_clean = []
+    messages_agora = []
+
+    for message in messages_raw:
+        role = message['role']
+        message_without_role = dict(message)
+        del message_without_role['role']
+
+        messages_agora.append({
+            'role': role,
+            'content': '```\n' + json.dumps(message_without_role, indent=2) + '\n```'
+        })
+
+        if message.get('status') == 'error':
+            messages_clean.append({
+                'role': role,
+                'content': f"Error: {message['message']}"
+            })
+        else:
+            messages_clean.append({
+                'role': role,
+                'content': message['body']
+            })
+
+    return messages_clean, messages_agora
+
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("### Agora Demo")
+        gr.Markdown("We will create a new Agora channel and offer it to Alice as a tool.")
+
+        chosen_task = gr.Dropdown(choices=list(SCHEMAS.keys()), label="Schema", value="weather_forecast")
+        custom_task = gr.Checkbox(label="Custom Task")
+
+        STATE_TRACKER = {}
+        
+        @gr.render(inputs=[chosen_task, custom_task])
+        def render(chosen_task, custom_task):
+            if STATE_TRACKER.get('chosen_task') != chosen_task:
+                STATE_TRACKER['chosen_task'] = chosen_task
+                for k, v in SCHEMAS[chosen_task].items():
+                    if isinstance(v, str):
+                        STATE_TRACKER[k] = v
+                    else:
+                        STATE_TRACKER[k] = json.dumps(v, indent=2)
+
+            if custom_task:
+                gr.Text(label="Description", value=STATE_TRACKER["description"], interactive=True).change(lambda x: STATE_TRACKER.update({'description': x}))
+                gr.TextArea(label="Input Schema", value=STATE_TRACKER["input"], interactive=True).change(lambda x: STATE_TRACKER.update({'input': x}))
+                gr.TextArea(label="Output Schema", value=STATE_TRACKER["output"], interactive=True).change(lambda x: STATE_TRACKER.update({'output': x}))
+                gr.TextArea(label="Tools", value=STATE_TRACKER["tools"], interactive=True).change(lambda x: STATE_TRACKER.update({'tools': x}))
+                gr.TextArea(label="Examples", value=STATE_TRACKER["examples"], interactive=True).change(lambda x: STATE_TRACKER.update({'examples': x}))
+
+        model_options = [
+            ('GPT 4o (Camel AI)', 'gpt-4o'),
+            ('GPT 4o-mini (Camel AI)', 'gpt-4o-mini'),
+            ('Claude 3 Sonnet (LangChain)', 'claude-3-sonnet'),
+            ('Gemini 1.5 Pro (Google GenAI)', 'gemini-1.5-pro'),
+            ('Llama3 405B (Sambanova + LangChain)', 'llama3-405b')
+        ]
+
+        fallback_image = ''
+
+        images = {
+            'gpt-4o': 'https://uxwing.com/wp-content/themes/uxwing/download/brands-and-social-media/chatgpt-icon.png',
+            'gpt-4o-mini': 'https://uxwing.com/wp-content/themes/uxwing/download/brands-and-social-media/chatgpt-icon.png',
+            'claude-3-sonnet': 'https://play-lh.googleusercontent.com/4S1nfdKsH_1tJodkHrBHimqlCTE6qx6z22zpMyPaMc_Rlr1EdSFDI1I6UEVMnokG5zI',
+            'gemini-1.5-pro': 'https://uxwing.com/wp-content/themes/uxwing/download/brands-and-social-media/google-gemini-icon.png',
+            'llama3-405b': 'https://www.designstub.com/png-resources/wp-content/uploads/2023/03/meta-icon-social-media-flat-graphic-vector-3-novem.png'
+        }
+
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
+                alice_model_dd = gr.Dropdown(label="Alice Model", choices=model_options, value="gpt-4o")
+            with gr.Column(scale=1):
+                bob_model_dd = gr.Dropdown(label="Bob Model", choices=model_options, value="gpt-4o")
+        
+        button = gr.Button('Start', elem_id='start_button')
+        gr.Markdown('### Natural Language')
+
+        @gr.render(inputs=[alice_model_dd, bob_model_dd])
+        def render_with_images(alice_model, bob_model):
+            avatar_images = [images.get(alice_model, fallback_image), images.get(bob_model, fallback_image)]
+            chatbot_nl = gr.Chatbot(type="messages", avatar_images=avatar_images)
+
+            with gr.Accordion(label="Raw Messages", open=False):
+                chatbot_nl_raw = gr.Chatbot(type="messages", avatar_images=avatar_images)
+        
+            gr.Markdown('### Negotiation')
+            chatbot_negotiation = gr.Chatbot(type="messages", avatar_images=avatar_images)
+
+            gr.Markdown('### Protocol')
+            protocol_result = gr.TextArea(interactive=False, label="Protocol")
+
+            gr.Markdown('### Implementation')
+            with gr.Row():
+                with gr.Column(scale=1):
+                    alice_implementation = gr.TextArea(interactive=False, label="Alice Implementation")
+                with gr.Column(scale=1):
+                    bob_implementation = gr.TextArea(interactive=False, label="Bob Implementation")
+
+            gr.Markdown('### Structured Communication')
+            structured_communication = gr.Chatbot(type="messages", avatar_images=avatar_images)
+            
+            with gr.Accordion(label="Raw Messages", open=False):
+                structured_communication_raw = gr.Chatbot(type="messages", avatar_images=avatar_images)
+
+            def respond(chosen_task, custom_task, alice_model, bob_model):
+                yield gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), \
+                    None, None, None, None, None, None, None, None
+
+                if custom_task:
+                    schema = dict(STATE_TRACKER)
+                    for k, v in schema.items():
+                        if isinstance(v, str):
+                            try:
+                                schema[k] = json.loads(v)
+                            except:
+                                pass
+                else:
+                    schema = SCHEMAS[chosen_task]
+
+                for nl_messages_raw, negotiation_messages, structured_messages_raw, protocol, alice_implementation, bob_implementation in full_flow(schema, alice_model, bob_model):
+                    nl_messages_clean, nl_messages_agora = parse_raw_messages(nl_messages_raw)
+                    structured_messages_clean, structured_messages_agora = parse_raw_messages(structured_messages_raw)
+
+                    yield gr.update(), gr.update(), gr.update(), nl_messages_clean, nl_messages_agora, negotiation_messages, structured_messages_clean, structured_messages_agora, protocol, alice_implementation, bob_implementation
+
+                #yield from full_flow(schema, alice_model, bob_model)
+                yield gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+
+            button.click(respond, [chosen_task, custom_task, alice_model_dd, bob_model_dd], [button, alice_model_dd, bob_model_dd, chatbot_nl, chatbot_nl_raw, chatbot_negotiation, structured_communication, structured_communication_raw, protocol_result, alice_implementation, bob_implementation])
+
+    demo.launch()
+
+
+if __name__ == '__main__':
+    main()

executor.py

@@ -0,0 +1,25 @@
+from abc import abstractmethod
+import importlib
+from typing import List
+
+from toolformers.base import Tool
+
+class Executor:
+    @abstractmethod
+    def run_routine(self, protocol_id, task_data, tools):
+        pass
+
+class UnsafeExecutor(Executor):
+    def run_routine(self, protocol_id, code, task_data, tools : List[Tool]):
+        protocol_id = protocol_id.replace('-', '_').replace('.', '_').replace('/', '_')
+        # TODO: This should be done in a safe, containerized environment
+        spec = importlib.util.spec_from_loader(protocol_id, loader=None)
+        loaded_module = importlib.util.module_from_spec(spec)
+
+        #spec.loader.exec_module(loaded_module)
+        exec(code, loaded_module.__dict__)
+
+        for tool in tools:
+            loaded_module.__dict__[tool.name] = tool.as_executable_function()
+
+        return loaded_module.run(task_data)

flow.py

@@ -0,0 +1,183 @@
+import dotenv
+dotenv.load_dotenv()
+
+import json
+import os
+import random
+import threading
+import time
+
+from toolformers.base import Tool, parameter_from_openai_api, StringParameter
+from toolformers.base import Toolformer
+from toolformers.camel import make_openai_toolformer
+from toolformers.langchain_agent import LangChainAnthropicToolformer
+from toolformers.sambanova import SambanovaToolformer
+from toolformers.gemini import GeminiToolformer
+
+from querier import Querier
+from responder import Responder
+from negotiator import SenderNegotiator, ReceiverNegotiator
+from programmer import SenderProgrammer, ReceiverProgrammer
+from executor import UnsafeExecutor
+from utils import compute_hash
+
+
+def create_toolformer(model_name) -> Toolformer:
+    if model_name in ['gpt-4o', 'gpt-4o-mini']:
+        return make_openai_toolformer(model_name)
+    elif 'claude' in model_name:
+        return LangChainAnthropicToolformer(model_name, os.environ.get('ANTHROPIC_API_KEY'))
+    elif model_name in ['llama3-405b']:
+        return SambanovaToolformer(model_name)
+    elif model_name in ['gemini-1.5-pro']:
+        return GeminiToolformer(model_name)
+    else:
+        raise ValueError(f"Unknown model name: {model_name}")
+
+def full_flow(schema, alice_model, bob_model):
+    NL_MESSAGES = []
+    NEGOTIATION_MESSAGES = []
+    STRUCTURED_MESSAGES = []
+    ARTIFACTS = {}
+
+    toolformer_alice = create_toolformer(alice_model)
+    toolformer_bob = create_toolformer(bob_model)
+
+    querier = Querier(toolformer_alice)
+    responder = Responder(toolformer_bob)
+
+    tools = []
+
+    for tool_schema in schema['tools']:
+        parameters = [parameter_from_openai_api(name, schema, name in tool_schema['input']['required']) for name, schema in tool_schema['input']['properties'].items()]
+        
+        def tool_fn(*args, **kwargs):
+            print(f'Bob tool {tool_schema["name"]} called with args {args} and kwargs {kwargs}')
+            return random.choice(tool_schema['dummy_outputs'])
+        
+        tool = Tool(tool_schema['name'], tool_schema['description'], parameters, tool_fn, tool_schema['output'])
+        tools.append(tool)
+
+    def nl_callback_fn(query):
+        print(query)
+        NL_MESSAGES.append({
+            'role': 'assistant',
+            #'content': query['body'],
+            'body': query['body'],
+            'protocolHash': None
+        })
+
+        response = responder.reply_to_query(query['body'], query['protocolHash'], tools, '')
+
+        NL_MESSAGES.append({
+            'role': 'user',
+            #'content': response['body']
+            'status': 'success',
+            'body': response['body']
+        })
+
+        return response
+    
+    negotiator_sender = SenderNegotiator(toolformer_alice)
+    negotiator_receiver = ReceiverNegotiator(toolformer_bob, tools, '')
+
+    def negotiation_callback_fn(query):
+        print(query)
+        NEGOTIATION_MESSAGES.append({
+            'role': 'assistant',
+            'content': query
+        })
+
+        response = negotiator_receiver.handle_negotiation(query)
+
+        NEGOTIATION_MESSAGES.append({
+            'role': 'user',
+            'content': response
+        })
+
+        #print('CURRENT NEGOTIATION MESSAGES:', len(NEGOTIATION_MESSAGES))
+
+        return response
+    
+    def final_message_callback_fn(query):
+        NEGOTIATION_MESSAGES.append({
+            'role': 'assistant',
+            'content': query
+        })
+
+    sender_programmer = SenderProgrammer(toolformer_alice)
+    receiver_programmer = ReceiverProgrammer(toolformer_bob)
+
+    executor = UnsafeExecutor()
+
+    def structured_callback_fn(query):
+        STRUCTURED_MESSAGES.append({
+            'role': 'assistant',
+            #'content': query
+            'body': json.dumps(query) if isinstance(query, dict) else query,
+            'protocolHash': ARTIFACTS['protocol']['hash'],
+            'protocolSources': ['https://...']
+        })
+
+        try:
+            response = executor.run_routine(ARTIFACTS['protocol']['hash'], ARTIFACTS['implementation_receiver'], query, tools)
+        except Exception as e:
+            STRUCTURED_MESSAGES.append({
+                'role': 'user',
+                'status': 'error',
+                'message': str(e)
+            })
+            return 'Error'
+
+        STRUCTURED_MESSAGES.append({
+            'role': 'user',
+            #'content': response
+            'status': 'success',
+            'body': json.dumps(response) if isinstance(response, dict) else response
+        })
+
+        return response
+
+    def flow():
+        task_data = random.choice(schema['examples'])
+        querier.send_query_without_protocol(schema, task_data, nl_callback_fn)
+
+        #time.sleep(1)
+
+        res = negotiator_sender.negotiate_protocol_for_task(schema, negotiation_callback_fn, final_message_callback_fn=final_message_callback_fn)
+        protocol_hash = compute_hash(res['protocol'])
+        res['hash'] = protocol_hash
+        ARTIFACTS['protocol'] = res
+
+        protocol_document = res['protocol']
+
+        implementation_sender = sender_programmer.write_routine_for_task(schema, protocol_document)
+        
+        ARTIFACTS['implementation_sender'] = implementation_sender
+
+        implementation_receiver = receiver_programmer.write_routine_for_tools(tools, protocol_document, '')
+
+        ARTIFACTS['implementation_receiver'] = implementation_receiver
+        send_tool = Tool('send_to_server', 'Send to server', StringParameter('query', 'The query', True), structured_callback_fn)
+
+        try:
+            executor.run_routine(protocol_hash, implementation_sender, task_data, [send_tool])
+        except Exception as e:
+            STRUCTURED_MESSAGES.append({
+                'role': 'assistant',
+                'status': 'success',
+                'error': str(e)
+            })
+
+    def get_info():
+        return NL_MESSAGES, NEGOTIATION_MESSAGES, STRUCTURED_MESSAGES, ARTIFACTS.get('protocol', {}).get('protocol', ''), \
+            ARTIFACTS.get('implementation_sender', ''), ARTIFACTS.get('implementation_receiver', '')
+
+    thread = threading.Thread(
+        target = lambda: flow()
+    )
+    thread.start()
+    while thread.is_alive():
+        yield get_info()
+        time.sleep(0.2)
+    yield get_info()

negotiator.py

@@ -0,0 +1,116 @@
+import json
+import uuid
+
+from toolformers.base import Toolformer
+from utils import extract_substring
+
+NEGOTIATION_RULES = '''
+Here are some rules (that should also be explained to the other GPT):
+- You can assume that the protocol has a sender and a receiver. Do not worry about how the messages will be delivered, focus only on the content of the messages.
+- Keep the protocol short and simple. It should be easy to understand and implement.
+- The protocol must specify the exact format of what is sent and received. Do not leave it open to interpretation.
+- The implementation will be written by a programmer that does not have access to the negotiation process, so make sure the protocol is clear and unambiguous.
+- The implementation will receive a string and return a string, so structure your protocol accordingly.
+- The other party might have a different internal data schema or set of tools, so make sure that the protocol is flexible enough to accommodate that.
+- There will only be one message sent by the sender and one message sent by the receiver. Design the protocol accordingly.
+- Keep the negotiation short: no need to repeat the same things over and over.
+- If the other party has proposed a protocol and you're good with it, there's no reason to keep negotiating or to repeat the protocol to the other party.
+- Do not restate parts of the protocols that have already been agreed upon.
+And remember: keep the protocol as simple and unequivocal as necessary. The programmer that will implement the protocol can code, but they are not a mind reader.
+'''
+
+TASK_NEGOTIATOR_PROMPT = f'''
+You are ProtocolNegotiatorGPT. Your task is to negotiate a protocol that can be used to query a service.
+You will receive a JSON schema of the task that the service must perform. Negotiate with the service to determine a protocol that can be used to query it.
+To do so, you will chat with another GPT (role: user) that will negotiate on behalf of the service.
+{NEGOTIATION_RULES}
+Once you are ready to save the protocol, reply wrapping the final version of the protocol, as agreed in your negotiation, between the tags <FINALPROTOCOL> and </FINALPROTOCOL>.
+Within the body of the tag, add the tags <NAME></NAME> and <DESCRIPTION></DESCRIPTION> to specify the name and description of the protocol.
+
+Remember that the <FINALPROTOCOL></FINALPROTOCOL> tags should also contain the protocol itself. Nothing outside such tags will be stored.
+'''
+
+class SenderNegotiator:
+    def __init__(self, toolformer : Toolformer):
+        self.toolformer = toolformer
+
+    def negotiate_protocol_for_task(self, task_schema, callback_fn, final_message_callback_fn=None):
+        found_protocol = None
+
+        prompt = TASK_NEGOTIATOR_PROMPT + '\nThe JSON schema of the task is the following:\n\n' + json.dumps(task_schema, indent=2)
+
+        conversation = self.toolformer.new_conversation(prompt, [], category='negotiation')
+
+        other_message = 'Hello! How may I help you?'
+        conversation_id = None
+
+        for i in range(10):
+            print('===NegotiatorGPT===')
+            message = conversation.chat(other_message, print_output=True)
+
+            print('Checking if we can extract from:', message)
+            print('---------')
+            protocol = extract_substring(message, '<FINALPROTOCOL>', '</FINALPROTOCOL>')
+
+            if protocol is None:
+                print('Could not extract')
+                other_message = callback_fn(message)
+                print()
+                print('===Other GPT===')
+                print(other_message)
+                print()
+            else:
+                if final_message_callback_fn:
+                    rest_of_message = message.split('<FINALPROTOCOL>')[0]
+                    final_message_callback_fn(rest_of_message)
+
+                name = extract_substring(protocol, '<NAME>', '</NAME>')
+                description = extract_substring(protocol, '<DESCRIPTION>', '</DESCRIPTION>')
+
+                if name is None:
+                    name = 'Unnamed protocol'
+                if description is None:
+                    description = 'No description provided'
+                
+                found_protocol = {
+                    'name': name,
+                    'description': description,
+                    'protocol': protocol
+                }
+                break
+
+        return found_protocol
+    
+TOOLS_NEGOTIATOR_PROMPT = f'''
+You are ProtocolNegotiatorGPT. You are negotiating a protocol on behalf of a web service that can perform a task.
+The other party is a GPT that is negotiating on behalf of the user. Your goal is to negotiate a protocol that is simple and clear, \
+but also expressive enough to allow the service to perform the task. A protocol is sufficiently expressive if you could write code \
+that, given the query formatted according to the protocol and the tools at the service's disposal, can parse the query according to \
+the protocol's specification, perform the task (if any) and send a reply.
+{NEGOTIATION_RULES}
+You will receive a list of tools that are available to the programmer that will implement the protocol.
+When you are okay with the protocol, don't further repeat everything, just tell to the other party that you are done.
+'''
+
+class ReceiverNegotiator:
+    def __init__(self, toolformer : Toolformer, tools, additional_info):
+        prompt = TOOLS_NEGOTIATOR_PROMPT
+
+        prompt += '\n\n' + additional_info
+
+        prompt += '\n\nThe tools that the implementer will have access to are:\n\n'
+
+        if len(tools) == 0:
+            prompt += 'No additional tools provided'
+        else:
+            for tool in tools:
+                prompt += tool.as_documented_python() + '\n\n'
+
+        print('Prompt:', prompt)
+
+        self.conversation = toolformer.new_conversation(prompt, tools, category='negotiation')
+
+    def handle_negotiation(self, message):
+        reply = self.conversation.chat(message, print_output=True)
+
+        return reply

programmer.py

@@ -0,0 +1,125 @@
+# The programmer creates implementations depending on a protocol specification
+
+import json
+import os
+
+from toolformers.base import Toolformer
+from utils import extract_substring
+
+TASK_PROGRAMMER_PROMPT = '''
+You are ProtocolProgrammerGPT. You will act as an intermediate between a machine (that has a certain input and output schema in JSON) \
+and a remote server that can perform a task following a certain protocol. Your task is to write a routine that takes some task data \
+(which follows the input schema), sends query in a format defined by the protocol, parses it and returns the output according to the output schema so that \
+the machine can use it.
+The routine is a Python file that contains a function "send_query". send_query takes a single argument, "task_data", which is a dictionary, and must return \
+a dictionary, which is the response to the query formatted according to the output schema.
+In order to communicate with the remote server, you can use the function "send_to_server" that is already available in the environment.
+send_to_server takes a single argument, "query" (which is a string formatted according to the protocol), and returns a string (again formatted according \
+to the protocol). Do not worry about managing communication, everything is already set up for you. Just focus on preparing the right query.
+
+Rules:
+- The implementation must be written in Python.
+- You can define any number of helper functions and import any libraries that are part of the Python standard library.
+- Do not import libraries that are not part of the Python standard library.
+- send_to_server will be already available in the environment. There is no need to import it.
+- Your task is to prepare the query, send it and parse the response.
+- Remember to import standard libraries if you need them.
+- If there is an unexpected error that is not covered by the protocol, throw an exception.\
+ If instead the protocol specifies how to handle the error, return the response according to the protocol's specification.
+- Do not execute anything (aside from library imports) when the file itself is loaded. I will personally import the file and call the send_query function with the task data.
+Begin by thinking about the implementation and how you would structure the code. \
+Then, write your implementation by writing a code block that contains the tags <IMPLEMENTATION> and </IMPLEMENTATION>. For example:
+```python
+<IMPLEMENTATION>
+
+def send_query(task_data):
+  ...
+
+</IMPLEMENTATION>
+'''
+
+class SenderProgrammer:
+    def __init__(self, toolformer : Toolformer):
+        self.toolformer = toolformer
+
+    def write_routine_for_task(self, task_schema, protocol_document):
+        conversation = self.toolformer.new_conversation(TASK_PROGRAMMER_PROMPT, [], category='programming')
+        message = 'JSON schema:\n\n' + json.dumps(task_schema) + '\n\n' + 'Protocol document:\n\n' + protocol_document
+
+        for i in range(5):
+            reply = conversation.chat(message, print_output=True)
+
+            implementation = extract_substring(reply, '<IMPLEMENTATION>', '</IMPLEMENTATION>')
+
+            if implementation is not None:
+                break
+
+            message = 'You have not provided an implementation yet. Please provide one by surrounding it in the tags <IMPLEMENTATION> and </IMPLEMENTATION>.'
+
+        implementation = implementation.strip()
+
+        # Sometimes the LLM leaves the Markdown formatting in the implementation
+        implementation = implementation.replace('```python', '').replace('```', '').strip()
+
+        implementation = implementation.replace('def send_query(', 'def run(')
+
+        return implementation
+
+TOOL_PROGRAMMER_PROMPT = '''
+You are ProtocolProgrammerGPT. Your task is to write a routine that takes a query formatted according to the protocol and returns a response.
+The routine is a Python file that contains a function "reply". reply takes a single argument, "query", which is a string, and must return a string.
+Depending on the protocol, the routine might be need to perform some actions before returning the response. The user might provide you with a list of \
+Python functions you can call to help you with this task. You don't need to worry about importing them, they are already available in the environment.
+Rules:
+- The implementation must be written in Python.
+- You can define any number of helper functions and import any libraries that are part of the Python standard library.
+- Do not import libraries that are not part of the Python standard library.
+- Remember to import standard libraries if you need them.
+- If there is an unexpected error that is not covered by the protocol, throw an exception.\
+ If instead the protocol specifies how to handle the error, return the response according to the protocol's specification.
+- Do not execute anything (aside from library imports) when the file itself is loaded. I will personally import the file and call the reply function with the task data.
+Begin by thinking about the implementation and how you would structure the code. \
+Then, write your implementation by writing a code block that contains the tags <IMPLEMENTATION> and </IMPLEMENTATION>. For example:
+```python
+<IMPLEMENTATION>
+
+def reply(query):
+  ...
+
+</IMPLEMENTATION>
+'''
+
+class ReceiverProgrammer:
+    def __init__(self, toolformer : Toolformer):
+        self.toolformer = toolformer
+    
+    def write_routine_for_tools(self, tools, protocol_document, additional_info):
+        conversation = self.toolformer.new_conversation(TOOL_PROGRAMMER_PROMPT + additional_info, [], category='programming')
+
+        message = 'Protocol document:\n\n' + protocol_document + '\n\n' + 'Additional functions:\n\n'
+
+        if len(tools) == 0:
+            message += 'No additional functions provided'
+        else:
+            for tool in tools:
+                message += tool.as_documented_python() + '\n\n'
+
+
+        for i in range(5):
+            reply = conversation.chat(message, print_output=True)
+
+            implementation = extract_substring(reply, '<IMPLEMENTATION>', '</IMPLEMENTATION>')
+
+            if implementation is not None:
+                break
+
+            message = 'You have not provided an implementation yet. Please provide one by surrounding it in the tags <IMPLEMENTATION> and </IMPLEMENTATION>.'
+
+        implementation = implementation.strip()
+
+        # Sometimes the LLM leaves the Markdown formatting in the implementation
+        implementation = implementation.replace('```python', '').replace('```', '').strip()
+
+        implementation = implementation.replace('def reply(', 'def run(')
+
+        return implementation

querier.py

@@ -0,0 +1,158 @@
+# The querier queries a service based on a protocol document.
+# It receives the protocol document and writes the query that must be performed to the system.
+
+import json
+import os
+from pathlib import Path
+
+import sys
+
+from toolformers.base import Tool, StringParameter, parameter_from_openai_api
+
+#from utils import send_raw_query
+
+PROTOCOL_QUERIER_PROMPT = 'You are QuerierGPT. You will receive a protocol document detailing how to query a service. Reply with a structured query which can be sent to the service.' \
+    'Only reply with the query itself, with no additional information or escaping. Similarly, do not add any additional whitespace or formatting.'
+
+def construct_query_description(protocol_document, task_schema, task_data):
+    query_description = ''
+    if protocol_document is not None:
+        query_description += 'Protocol document:\n\n'
+        query_description += protocol_document + '\n\n'
+    query_description += 'JSON schema of the task:\n\n'
+    query_description += 'Input (i.e. what the machine will provide you):\n'
+    query_description += json.dumps(task_schema['input'], indent=2) + '\n\n'
+    query_description += 'Output (i.e. what you have to provide to the machine):\n'
+    query_description += json.dumps(task_schema['output'], indent=2) + '\n\n'
+    query_description += 'JSON data of the task:\n\n'
+    query_description += json.dumps(task_data, indent=2) + '\n\n'
+
+    return query_description
+
+NL_QUERIER_PROMPT = 'You are NaturalLanguageQuerierGPT. You act as an intermediary between a machine (who has a very specific input and output schema) and an agent (who uses natural language).' \
+    'You will receive a task description (including a schema of the input and output) that the machine uses and the corresponding data. Call the \"sendQuery\" tool with a natural language message where you ask to perform the task according to the data.' \
+    'Make sure to mention all the relevant information. ' \
+    'Do not worry about managing communication, everything is already set up for you. Just focus on asking the right question.' \
+    'The sendQuery tool will return the reply of the service.\n' \
+    'Once you receive the reply, call the \"deliverStructuredOutput\" tool with parameters according to the task\'s output schema. \n' \
+    'Note: you cannot call sendQuery multiple times, so make sure to ask the right question the first time. Similarly, you cannot call deliverStructuredOutput multiple times, so make sure to deliver the right output the first time.' \
+    'If the query fails, do not attempt to send another query.'
+
+def parse_and_handle_query(query, callback_fn, protocol_id, source):
+    if isinstance(query, dict):
+        query = json.dumps(query)
+    response = callback_fn({
+        "protocolHash": protocol_id,
+        "body": query,
+        "protocolSources": None if source is None else [source]
+    })
+
+    print('Response:', response, type(response))
+
+    if response['status'] == 'success':
+        return response['body']
+    else:
+        return 'Error calling the tool: ' + response['message']
+
+def get_output_parameters(task_schema):
+    output_schema = task_schema['output']
+    required_parameters = output_schema['required']
+
+    parameters = []
+
+    for parameter_name, parameter_schema in output_schema['properties'].items():
+        parameter = parameter_from_openai_api(parameter_name, parameter_schema, parameter_name in required_parameters)
+        parameters.append(parameter)
+    
+    return parameters
+
+class Querier:
+    def __init__(self, toolformer):
+        self.toolformer = toolformer
+
+    def handle_conversation(self, prompt, message, callback_fn, protocol_id, source, output_parameters):
+        sent_query_counter = 0
+        
+        def send_query_internal(query):
+            print('Sending query:', query)
+            nonlocal sent_query_counter
+            sent_query_counter += 1
+
+            if sent_query_counter > 50:
+                # All hope is lost, crash
+                sys.exit(-2)
+            elif sent_query_counter > 10:
+                # LLM is not listening, throw an exception
+                raise Exception('Too many attempts to send queries. Exiting.')
+            elif sent_query_counter > 5:
+                # LLM is not listening, issue a warning
+                return 'You have attempted to send too many queries. Finish the message and allow the user to speak, or the system will crash.'
+            elif sent_query_counter > 1:
+                return 'You have already sent a query. You cannot send another one.'
+            return parse_and_handle_query(query, callback_fn, protocol_id, source)
+
+        send_query_tool = Tool('sendQuery', 'Send a query to the other service based on a protocol document.', [
+            StringParameter('query', 'The query to send to the service', True)
+        ], send_query_internal)
+
+        found_output = None
+        registered_output_counter = 0
+
+        def register_output(**kwargs):
+            print('Registering output:', kwargs)
+
+            nonlocal found_output
+            nonlocal registered_output_counter
+            if found_output is not None:
+                registered_output_counter += 1
+
+            if registered_output_counter > 50:
+                # All hope is lost, crash
+                sys.exit(-2)
+            elif registered_output_counter > 10:
+                # LLM is not listening, raise an exception
+                raise Exception('Too many attempts to register outputs. Exiting.')
+            elif registered_output_counter > 5:
+                # LLM is not listening, issue a warning
+                return 'You have attempted to register too many outputs. Finish the message and allow the user to speak, or the system will crash.'
+            elif registered_output_counter > 0:
+                return 'You have already registered an output. You cannot register another one.'
+
+            output = json.dumps(kwargs)
+
+            found_output = output
+            return 'Done'
+
+        register_output_tool = Tool('deliverStructuredOutput', 'Deliver the structured output to the machine.',
+            output_parameters
+        , register_output)
+
+        conversation = self.toolformer.new_conversation(prompt, [send_query_tool, register_output_tool], category='conversation')
+
+        for i in range(5):
+            conversation.chat(message, print_output=True)
+
+            if found_output is not None:
+                break
+
+            # If we haven't sent a query yet, we can't proceed
+            if sent_query_counter == 0:
+                message = 'You must send a query before delivering the structured output.'
+            elif found_output is None:
+                message = 'You must deliver the structured output.'
+
+        return found_output
+
+    #def send_query_with_protocol(self, storage, task_schema, task_data, target_node, protocol_id, source):
+    #    base_folder = Path(os.environ.get('STORAGE_PATH')) / 'protocol_documents'
+    #    protocol_document = storage.load_protocol_document(base_folder, protocol_id)
+    #    query_description = construct_query_description(protocol_document, task_schema, task_data)
+    #    output_parameters = get_output_parameters(task_schema)
+    #
+    #    return self.handle_conversation(PROTOCOL_QUERIER_PROMPT, query_description, target_node, protocol_id, source, output_parameters)
+
+    def send_query_without_protocol(self, task_schema, task_data, callback_fn):
+        query_description = construct_query_description(None, task_schema, task_data)
+        output_parameters = get_output_parameters(task_schema)
+
+        return self.handle_conversation(NL_QUERIER_PROMPT, query_description, callback_fn, None, None, output_parameters)

requirements.txt

@@ -0,0 +1,28 @@
+asknews==0.7.51
+camel-ai==0.2.6
+google-ai-generativelanguage==0.6.4
+google-api-core==2.22.0
+google-api-python-client==2.151.0
+google-auth==2.36.0
+google-auth-httplib2==0.2.0
+google-cloud-core==2.4.1
+google-cloud-storage==2.18.2
+google-cloud-vision==3.8.0
+google-crc32c==1.6.0
+google-generativeai==0.6.0
+gradio==5.5.0
+gradio_client==1.4.2
+huggingface-hub==0.26.2
+langchain==0.3.7
+langchain-anthropic==0.2.4
+langchain-community==0.3.5
+langchain-core==0.3.15
+langchain-text-splitters==0.3.2
+langgraph==0.2.45
+langgraph-checkpoint==2.0.2
+langgraph-sdk==0.1.35
+python-dotenv==1.0.1
+requests-oauthlib==1.3.1
+sseclient-py==1.8.0
+torch==2.5.1
+transformers==4.46.2

responder.py

@@ -0,0 +1,87 @@
+# The responder is a special toolformer that replies to a service based on a protocol document.
+# It receives the protocol document and writes the response that must be sent to the system.
+
+import json
+import os
+from pathlib import Path
+
+from toolformers.base import Toolformer
+
+# TODO: A tool to declare an error?
+
+
+PROTOCOL_RESPONDER_PROMPT = 'You are ResponderGPT. You will receive a protocol document detailing how to respond to a query. '\
+    'Use the provided functions to execute what is requested and provide the response according to the protocol\'s specification. ' \
+    'Only reply with the response itself, with no additional information or escaping. Similarly, do not add any additional whitespace or formatting.'# \
+   # 'If you do not have enough information to reply, or if you cannot execute the request, reply with "ERROR" (without quotes).'
+
+NL_RESPONDER_PROMPT = 'You are NaturalLanguageResponderGPT. You will receive a query from a user. ' \
+    'Use the provided functions to execute what is requested and reply with a response (in natural language). ' \
+    'Important: the user does not have the capacity to respond to follow-up questions, so if you think you have enough information to reply/execute the actions, do so.'
+    #'If you do not have enough information to reply, if you cannot execute the request, or if the request is invalid, reply with "ERROR" (without quotes).' \
+
+
+    
+class Responder:
+    def __init__(self, toolformer : Toolformer):
+        self.toolformer = toolformer
+
+    def reply_with_protocol_document(self, query, protocol_document, tools, additional_info):
+        print('===NL RESPONDER (WITH PROTOCOL)===')
+
+        conversation = self.toolformer.new_conversation(PROTOCOL_RESPONDER_PROMPT + additional_info, tools, category='conversation')
+
+        prompt = 'The protocol is the following:\n\n' + protocol_document + '\n\nThe query is the following:' + query
+
+        reply = conversation.chat(prompt, print_output=True)
+
+        print('======')
+
+        if 'error' in reply.lower().strip()[-10:]:
+            return {
+                'status': 'error',
+                'message': 'Error in the response'
+            }
+
+        return {
+            'status': 'success',
+            'body': reply
+        }
+
+
+    def reply_to_nl_query(self, query, tools, additional_info):
+        print('===NL RESPONDER (NO PROTOCOL)===')
+        print(NL_RESPONDER_PROMPT + additional_info)
+        print([tool.name for tool in tools])
+
+        conversation = self.toolformer.new_conversation(NL_RESPONDER_PROMPT + additional_info, tools, category='conversation')
+
+        print('Created conversation')
+        try:
+            reply = conversation.chat(query, print_output=True)
+        except Exception as e:
+            # Print traceback
+            import traceback
+            traceback.print_exc()
+            raise e
+        print('======')
+
+        if 'error' in reply.lower().strip()[-10:]:
+            return {
+                'status': 'error',
+            }
+
+        return {
+            'status': 'success',
+            'body': reply
+        }
+
+
+    def reply_to_query(self, query, protocol_id, tools, additional_info):
+        print('Additional info:', additional_info)
+        if protocol_id is None:
+            return self.reply_to_nl_query(query, tools, additional_info)
+        #else:
+        #    base_folder = Path(os.environ.get('STORAGE_PATH')) / 'protocol_documents'
+        #    protocol_document = self.memory.load_protocol_document(base_folder, protocol_id)
+        #    return self.reply_with_protocol_document(query, protocol_document, tools, additional_info)

toolformers/base.py

@@ -0,0 +1,332 @@
+from abc import ABC, abstractmethod
+
+import json
+
+from google.generativeai.types import CallableFunctionDeclaration
+import google.generativeai.types.content_types as content_types
+
+from utils import add_params_and_annotations
+
+class Parameter:
+    def __init__(self, name, description, required):
+        self.name = name
+        self.description = description
+        self.required = required
+
+    def as_openai_info(self):
+        pass
+
+    def as_standard_api(self):
+        pass
+
+class StringParameter(Parameter):
+    def __init__(self, name, description, required):
+        super().__init__(name, description, required)
+
+    def as_openai_info(self):
+        return {
+            "type": "string",
+            "name": self.name,
+            "description": self.description
+        }
+
+    def as_standard_api(self):
+        return {
+            "type": "string",
+            "name": self.name,
+            "description": self.description,
+            "required": self.required
+        }
+
+    def as_natural_language(self):
+        return f'{self.name} (string{", required" if self.required else ""}): {self.description}.'
+
+    def as_documented_python(self):
+        return f'{self.name} (str{", required" if self.required else ""}): {self.description}.'
+    
+    def as_gemini_tool(self):
+        return {
+            'type': 'string',
+            'description': self.description
+        }
+
+    @staticmethod
+    def from_standard_api(api_info):
+        return StringParameter(api_info["name"], api_info["description"], api_info["required"])
+
+class EnumParameter(Parameter):
+    def __init__(self, name, description, values, required):
+        super().__init__(name, description, required)
+        self.values = values
+
+    def as_openai_info(self):
+        return {
+            "type": "string",
+            "description": self.description,
+            "values": self.values
+        }
+
+    def as_standard_api(self):
+        return {
+            "type": "enum",
+            "name": self.name,
+            "description": self.description,
+            "values": self.values,
+            "required": self.required
+        }
+    
+    def as_natural_language(self):
+        return f'{self.name} (enum{", required" if self.required else ""}): {self.description}. Possible values: {", ".join(self.values)}'
+    
+    def as_documented_python(self):
+        return f'{self.name} (str{", required" if self.required else ""}): {self.description}. Possible values: {", ".join(self.values)}'
+
+    def as_gemini_tool(self):
+        return {
+            'description': self.description,
+            'type': 'string',
+            'enum': self.values
+        }
+    
+    @staticmethod
+    def from_standard_api(api_info):
+        return EnumParameter(api_info["name"], api_info["description"], api_info["values"], api_info["required"])
+
+class NumberParameter(Parameter):
+    def __init__(self, name, description, required):
+        super().__init__(name, description, required)
+
+    def as_openai_info(self):
+        return {
+            "type": "number",
+            "description": self.description
+        }
+
+    def as_standard_api(self):
+        return {
+            "type": "number",
+            "name": self.name,
+            "description": self.description,
+            "required": self.required
+        }
+
+    def as_natural_language(self):
+        return f'{self.name} (number): {self.description}'
+
+    def as_documented_python(self):
+        return f'{self.name} (number): {self.description}'
+    
+    def as_gemini_tool(self):
+        return {
+            'description': self.description,
+            'type': 'number'
+        }
+    
+class ArrayParameter(Parameter):
+    def __init__(self, name, description, required, item_schema):
+        super().__init__(name, description, required)
+        self.item_schema = item_schema
+
+    def as_openai_info(self):
+        return {
+            "type": "array",
+            "description": self.description,
+            "items": self.item_schema
+        }
+
+    def as_standard_api(self):
+        return {
+            "type": "array",
+            "name": self.name,
+            "description": self.description,
+            "required": self.required,
+            "item_schema": self.item_schema
+        }
+
+    def as_natural_language(self):
+        return f'{self.name} (array): {self.description}. Each item should follow the JSON schema: {json.dumps(self.item_schema)}'
+
+    def as_documented_python(self):
+        return f'{self.name} (list): {self.description}. Each item should follow the JSON schema: {json.dumps(self.item_schema)}'
+    
+    def as_gemini_tool(self):
+        return {
+            'description': self.description,
+            'type': 'array',
+            'items': self.item_schema
+        }
+
+def parameter_from_openai_api(parameter_name, schema, required):
+    if 'enum' in schema:
+        return EnumParameter(parameter_name, schema['description'], schema['enum'], required)
+    elif schema['type'] == 'string':
+        return StringParameter(parameter_name, schema['description'], required)
+    elif schema['type'] == 'number':
+        return NumberParameter(parameter_name, schema['description'], required)
+    elif schema['type'] == 'array':
+        return ArrayParameter(parameter_name, schema['description'], required, schema['items'])
+    else:
+        raise ValueError(f'Unknown parameter type: {schema["type"]}')
+
+class Tool:
+    def __init__(self, name, description, parameters, function, output_schema=None):
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+        self.function = function
+        self.output_schema = output_schema
+    
+    def call_tool_for_toolformer(self, *args, **kwargs):
+        print(f'Toolformer called tool {self.name} with args {args} and kwargs {kwargs}')
+        # Unlike a call from a routine, this call catches exceptions and returns them as strings
+        try:
+            tool_reply = self.function(*args, **kwargs)
+            print(f'Tool {self.name} returned: {tool_reply}')
+            return tool_reply
+        except Exception as e:
+            print(f'Tool {self.name} failed with exception: {e}')
+            return 'Tool call failed: ' + str(e)
+    
+    def as_openai_info(self):
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type" : "object",
+                    "properties": {parameter.name : parameter.as_openai_info() for parameter in self.parameters},
+                    "required": [parameter.name for parameter in self.parameters if parameter.required]
+                }
+            }
+        }
+    
+    def as_gemini_tool(self) -> CallableFunctionDeclaration:
+        if len(self.parameters) == 0:
+            parameters = None
+        else:
+            parameters = {
+                'type': 'object',
+                'properties': {parameter.name: parameter.as_gemini_tool() for parameter in self.parameters},
+                'required': [parameter.name for parameter in self.parameters if parameter.required]
+            }
+        return content_types.Tool([CallableFunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=parameters,
+            function=self.call_tool_for_toolformer
+        )])
+
+    def as_llama_schema(self):
+        schema = {
+            'name': self.name,
+            'description': self.description,
+            'parameters': {parameter.name : parameter.as_openai_info() for parameter in self.parameters},
+            'required': [parameter.name for parameter in self.parameters if parameter.required]
+        }
+
+        if self.output_schema is not None:
+            schema['output_schema'] = self.output_schema
+        
+        return schema
+
+    def as_natural_language(self):
+        print('Converting to natural language')
+        print('Number of parameters:', len(self.parameters))
+        nl = f'Function {self.name}: {self.description}. Parameters:\n'
+        
+        if len(self.parameters) == 0:
+            nl += 'No parameters.'
+        else:
+            for parameter in self.parameters:
+                nl += '\t' + parameter.as_natural_language() + '\n'
+
+        if self.output_schema is not None:
+            nl += f'\Returns a dictionary with schema: {json.dumps(self.output_schema, indent=2)}'
+        
+        return nl
+
+    def as_standard_api(self):
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": [parameter.as_standard_api() for parameter in self.parameters]
+        }
+    
+    def as_documented_python(self):
+        documented_python = f'Tool {self.name}:\n\n{self.description}\nParameters:\n'
+
+        if len(self.parameters) == 0:
+            documented_python += 'No parameters.'
+        else:
+            for parameter in self.parameters:
+                documented_python += '\t' + parameter.as_documented_python() + '\n'
+        
+        if self.output_schema is not None:
+            documented_python += f'\Returns a dictionary with schema: {json.dumps(self.output_schema, indent=2)}'
+
+        return documented_python
+
+    def as_executable_function(self):
+        # Create an actual function that can be called
+        def f(*args, **kwargs):
+            print('Routine called tool', self.name, 'with args', args, 'and kwargs', kwargs)
+            response = self.function(*args, **kwargs)
+            print('Tool', self.name, 'returned:', response)
+            return response
+        
+        return f
+    
+    def as_annotated_function(self):
+        def wrapped_fn(*args, **kwargs):
+            return self.call_tool_for_toolformer(*args, **kwargs)
+        
+        parsed_parameters = {}
+
+        description = self.description
+
+        for parameter_name, parameter_schema in self.as_openai_info()['function']['parameters']['properties'].items():
+            if parameter_schema['type'] == 'string':
+                parsed_parameters[parameter_name] = (str, parameter_schema['description'])
+            elif parameter_schema['type'] == 'number':
+                parsed_parameters[parameter_name] = (float, parameter_schema['description'])
+            elif parameter_schema['type'] == 'object':
+                parsed_parameters[parameter_name] = (dict, parameter_schema['description'])
+
+                description += f'\n{parameter_name} has the schema:\n' + json.dumps(parameter_schema) + '\n'
+            else:
+                raise ValueError(f'Unknown parameter type: {parameter_schema["type"]}')
+        
+        return_type = type(None)
+
+        if self.output_schema is not None:
+            #description += '\nOutput schema:\n' + json.dumps(self.output_schema)
+
+            if self.output_schema['type'] == 'string':
+                return_type = str
+            elif self.output_schema['type'] == 'number':
+                return_type = float
+            elif self.output_schema['type'] == 'object':
+                return_type = dict
+            else:
+                raise ValueError(f'Unknown output type: {self.output_schema["type"]}')
+
+        return add_params_and_annotations(
+            self.name, description, parsed_parameters, return_type)(wrapped_fn)
+    
+    @staticmethod
+    def from_openai_info(info, func):
+        parameters = [parameter_from_openai_api(name, schema, name in info['function']['parameters']['required']) for name, schema in info['function']['parameters']['properties'].items()]
+        return Tool(info['function']['name'], info['function']['description'], parameters, func)
+
+
+class Conversation(ABC):
+    @abstractmethod
+    def chat(self, message, role='user', print_output=True):
+        pass
+
+class Toolformer(ABC):
+    @abstractmethod
+    def new_conversation(self, prompt, tools, category=None) -> Conversation:
+        pass
+

toolformers/camel.py

@@ -0,0 +1,86 @@
+import datetime
+import os
+from typing import List
+import warnings
+
+from toolformers.base import Conversation, Toolformer, Tool
+from camel.messages import BaseMessage
+from camel.models import ModelFactory
+from camel.types import ModelPlatformType, ModelType
+from camel.messages import BaseMessage as bm
+from camel.agents import ChatAgent
+from camel.toolkits.function_tool import FunctionTool
+from camel.configs.openai_config import ChatGPTConfig
+
+class CamelConversation(Conversation):
+    def __init__(self, toolformer, agent, category=None):
+        self.toolformer = toolformer
+        self.agent = agent
+        self.category = category
+    
+    def chat(self, message, role='user', print_output=True):
+        agent_id = os.environ.get('AGENT_ID', None)
+
+        start_time = datetime.datetime.now()
+
+        if role == 'user':
+            formatted_message = BaseMessage.make_user_message('user', message)
+        elif role == 'assistant':
+            formatted_message = BaseMessage.make_assistant_message('assistant', message)
+        else:
+            raise ValueError('Role must be either "user" or "assistant".')
+        
+        response = self.agent.step(formatted_message)
+
+        reply = response.msg.content
+
+        if print_output:
+            print(reply)
+        
+        return reply
+
+class CamelToolformer(Toolformer):
+    def __init__(self, model_platform, model_type, model_config_dict, name=None):
+        self.model_platform = model_platform
+        self.model_type = model_type
+        self.model_config_dict = model_config_dict
+        self._name = name
+
+    @property
+    def name(self):
+        if self._name is None:
+            return f'{self.model_platform.value}_{self.model_type.value}'
+        else:
+            return self._name
+
+    def new_conversation(self, prompt, tools : List[Tool], category=None) -> Conversation:
+        model = ModelFactory.create(
+            model_platform=self.model_platform,
+            model_type=self.model_type,
+            model_config_dict=self.model_config_dict
+        )
+
+        agent = ChatAgent(
+            model=model,
+            system_message=bm.make_assistant_message('system', prompt),
+            tools=[FunctionTool(tool.call_tool_for_toolformer, openai_tool_schema=tool.as_openai_info()) for tool in tools]
+        )
+
+        return CamelConversation(self, agent, category)
+
+def make_openai_toolformer(model_type_internal):
+    if model_type_internal == 'gpt-4o':
+        model_type = ModelType.GPT_4O
+    elif model_type_internal == 'gpt-4o-mini':
+        model_type = ModelType.GPT_4O_MINI
+    else:
+        raise ValueError('Model type must be either "gpt-4o" or "gpt-4o-mini".')
+
+    #formatted_tools = [FunctionTool(tool.call_tool_for_toolformer, tool.as_openai_info()) for tool in tools]
+
+    return CamelToolformer(
+        model_platform=ModelPlatformType.OPENAI,
+        model_type=model_type,
+        model_config_dict=ChatGPTConfig(temperature=0.2).as_dict(),
+        name=model_type_internal
+    )

toolformers/gemini.py

@@ -0,0 +1,100 @@
+import datetime
+import os
+from random import random
+import time
+import traceback
+from typing import List
+
+from toolformers.base import Conversation, Tool, Toolformer
+
+import google.generativeai as genai
+from google.generativeai.generative_models import ChatSession
+
+genai.configure(api_key=os.environ['GOOGLE_API_KEY'])
+
+class GeminiConversation(Conversation):
+    def __init__(self, model_name, chat_agent : ChatSession, category=None):
+        self.model_name = model_name
+        self.chat_agent = chat_agent
+        self.category = category
+
+    def chat(self, message, role='user', print_output=True):
+        agent_id = os.environ.get('AGENT_ID', None)
+        time_start = datetime.datetime.now()
+
+        exponential_backoff_lower = 30
+        exponential_backoff_higher = 60
+        for i in range(5):
+            try:
+                response = self.chat_agent.send_message({
+                    'role': role,
+                    'parts': [
+                        message
+                    ]
+                })
+                break
+            except Exception as e:
+                print(e)
+                if '429' in str(e):
+                    print('Rate limit exceeded. Waiting with random exponential backoff.')
+                    if i < 4:
+                        time.sleep(random() * (exponential_backoff_higher - exponential_backoff_lower) + exponential_backoff_lower)
+                        exponential_backoff_lower *= 2
+                        exponential_backoff_higher *= 2
+                elif 'candidates[0]' in traceback.format_exc():
+                    # When Gemini has nothing to say, it raises an error with this message
+                    print('No response')
+                    return 'No response'
+                elif '500' in str(e):
+                    # Sometimes Gemini just decides to return a 500 error for absolutely no reason. Retry.
+                    print('500 error')
+                    time.sleep(5)
+                    traceback.print_exc()
+                else:
+                    raise e
+
+        time_end = datetime.datetime.now()
+
+        usage_info = {
+            'prompt_tokens': response.usage_metadata.prompt_token_count,
+            'completion_tokens': response.usage_metadata.candidates_token_count
+        }
+
+        #send_usage_to_db(
+        #    usage_info,
+        #    time_start,
+        #    time_end,
+        #    agent_id,
+        #    self.category,
+        #    self.model_name
+        #)
+
+        reply = response.text
+
+        if print_output:
+            print(reply)
+        
+        return reply
+
+class GeminiToolformer(Toolformer):
+    def __init__(self, model_name):
+        self.model_name = model_name
+
+    def new_conversation(self, system_prompt, tools : List[Tool], category=None) -> Conversation:
+        print('Tools:')
+        print('\n'.join([str(tool.as_openai_info()) for tool in tools]))
+        model = genai.GenerativeModel(
+            model_name=self.model_name,
+            system_instruction=system_prompt,
+            tools=[tool.as_gemini_tool() for tool in tools]
+        )
+
+        chat = model.start_chat(enable_automatic_function_calling=True)
+
+        return GeminiConversation(self.model_name, chat, category)
+
+def make_gemini_toolformer(model_name):
+    if model_name not in ['gemini-1.5-flash', 'gemini-1.5-pro']:
+        raise ValueError(f"Unknown model name: {model_name}")
+
+    return GeminiToolformer(model_name)

toolformers/huggingface_agent.py

@@ -0,0 +1,39 @@
+from huggingface_hub import InferenceClient
+from transformers import ReactCodeAgent
+from transformers import tool as function_to_tool
+
+from toolformers.base import Conversation, Toolformer
+
+class HuggingFaceConversation(Conversation):
+    def __init__(self, agent : ReactCodeAgent, prompt, category=None):
+        self.agent = agent
+        self.messages = [('system', prompt)]
+        self.category = category
+
+    def chat(self, message, role='user', print_output=True) -> str:
+        self.messages.append((role, message))
+        
+        final_prompt = 'For context, here are the previous messages in the conversation:\n\n'
+
+        for role, message in self.messages:
+            final_prompt += f'{role.capitalize()}: {message}\n'
+
+        final_prompt += "Don't worry, you don't need to use the same format to reply. Stick with the Task:/Action:/etc. format.\n\n"
+
+        response = self.agent.run(final_prompt)
+        print(response)
+        return response
+
+class HuggingFaceToolformer(Toolformer):
+    def __init__(self, model_name, max_tokens=2000):
+        self.model = InferenceClient(model=model_name)
+        self.max_tokens = max_tokens
+
+    def llm_engine(self, messages, stop_sequences=["Task"]) -> str:
+        response = self.model.chat_completion(messages, stop=stop_sequences, max_tokens=self.max_tokens)
+        answer = response.choices[0].message.content
+        return answer
+    
+    def new_conversation(self, prompt, tools, category=None):
+        agent = ReactCodeAgent(tools=[function_to_tool(tool.as_annotated_function()) for tool in tools], llm_engine=self.llm_engine)
+        return HuggingFaceConversation(agent, prompt, category=category)

toolformers/langchain_agent.py

@@ -0,0 +1,73 @@
+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
+
+
+
+class LangChainConversation(Conversation):
+    def __init__(self, agent, messages, category=None):
+        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 = ''
+        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
+            #final_message += chunk['agent']['messages'].content
+
+        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(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?'))

toolformers/sambanova/__init__.py

@@ -0,0 +1 @@
+from .core import SambanovaToolformer

toolformers/sambanova/api_gateway.py

@@ -0,0 +1,146 @@
+import logging
+import os
+import sys
+from typing import Optional
+
+from langchain_community.llms.sambanova import SambaStudio
+from langchain_core.language_models.llms import LLM
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+utils_dir = os.path.abspath(os.path.join(current_dir, '..'))
+repo_dir = os.path.abspath(os.path.join(utils_dir, '..'))
+sys.path.append(utils_dir)
+sys.path.append(repo_dir)
+
+from toolformers.sambanova.sambanova_langchain import SambaNovaCloud
+
+EMBEDDING_MODEL = 'intfloat/e5-large-v2'
+NORMALIZE_EMBEDDINGS = True
+
+# Configure the logger
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s [%(levelname)s] - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+    ],
+)
+logger = logging.getLogger(__name__)
+
+
+class APIGateway:
+    @staticmethod
+    def load_llm(
+        type: str,
+        streaming: bool = False,
+        coe: bool = False,
+        do_sample: Optional[bool] = None,
+        max_tokens_to_generate: Optional[int] = None,
+        temperature: Optional[float] = None,
+        select_expert: Optional[str] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        stop_sequences: Optional[str] = None,
+        process_prompt: Optional[bool] = False,
+        sambastudio_base_url: Optional[str] = None,
+        sambastudio_base_uri: Optional[str] = None,
+        sambastudio_project_id: Optional[str] = None,
+        sambastudio_endpoint_id: Optional[str] = None,
+        sambastudio_api_key: Optional[str] = None,
+        sambanova_url: Optional[str] = None,
+        sambanova_api_key: Optional[str] = None,
+    ) -> LLM:
+        """Loads a langchain Sambanova llm model given a type and parameters
+        Args:
+            type (str): wether to use sambastudio, or SambaNova Cloud model "sncloud"
+            streaming (bool): wether to use streaming method. Defaults to False.
+            coe (bool): whether to use coe model. Defaults to False.
+
+            do_sample (bool) : Optional wether to do sample.
+            max_tokens_to_generate (int) : Optional max number of tokens to generate.
+            temperature (float) : Optional model temperature.
+            select_expert (str) : Optional expert to use when using CoE models.
+            top_p (float) : Optional model top_p.
+            top_k (int) : Optional model top_k.
+            repetition_penalty (float) : Optional model repetition penalty.
+            stop_sequences (str) : Optional model stop sequences.
+            process_prompt (bool) : Optional default to false.
+
+            sambastudio_base_url (str): Optional SambaStudio environment URL".
+            sambastudio_base_uri (str): Optional SambaStudio-base-URI".
+            sambastudio_project_id (str): Optional SambaStudio project ID.
+            sambastudio_endpoint_id (str): Optional SambaStudio endpoint ID.
+            sambastudio_api_token (str): Optional SambaStudio endpoint API key.
+
+            sambanova_url (str): Optional SambaNova Cloud URL",
+            sambanova_api_key (str): Optional SambaNovaCloud API key.
+
+        Returns:
+            langchain llm model
+        """
+
+        if type == 'sambastudio':
+            envs = {
+                'sambastudio_base_url': sambastudio_base_url,
+                'sambastudio_base_uri': sambastudio_base_uri,
+                'sambastudio_project_id': sambastudio_project_id,
+                'sambastudio_endpoint_id': sambastudio_endpoint_id,
+                'sambastudio_api_key': sambastudio_api_key,
+            }
+            envs = {k: v for k, v in envs.items() if v is not None}
+            if coe:
+                model_kwargs = {
+                    'do_sample': do_sample,
+                    'max_tokens_to_generate': max_tokens_to_generate,
+                    'temperature': temperature,
+                    'select_expert': select_expert,
+                    'top_p': top_p,
+                    'top_k': top_k,
+                    'repetition_penalty': repetition_penalty,
+                    'stop_sequences': stop_sequences,
+                    'process_prompt': process_prompt,
+                }
+                model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None}
+
+                llm = SambaStudio(
+                    **envs,
+                    streaming=streaming,
+                    model_kwargs=model_kwargs,
+                )
+            else:
+                model_kwargs = {
+                    'do_sample': do_sample,
+                    'max_tokens_to_generate': max_tokens_to_generate,
+                    'temperature': temperature,
+                    'top_p': top_p,
+                    'top_k': top_k,
+                    'repetition_penalty': repetition_penalty,
+                    'stop_sequences': stop_sequences,
+                }
+                model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None}
+                llm = SambaStudio(
+                    **envs,
+                    streaming=streaming,
+                    model_kwargs=model_kwargs,
+                )
+
+        elif type == 'sncloud':
+            envs = {
+                'sambanova_url': sambanova_url,
+                'sambanova_api_key': sambanova_api_key,
+            }
+            envs = {k: v for k, v in envs.items() if v is not None}
+            llm = SambaNovaCloud(
+                **envs,
+                max_tokens=max_tokens_to_generate,
+                model=select_expert,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+            )
+
+        else:
+            raise ValueError(f"Invalid LLM API: {type}, only 'sncloud' and 'sambastudio' are supported.")
+
+        return llm

toolformers/sambanova/core.py

@@ -0,0 +1,46 @@
+import datetime
+import os
+from typing import List
+
+from toolformers.base import Conversation, Toolformer, Tool
+from toolformers.sambanova.function_calling import FunctionCallingLlm
+
+class SambanovaConversation(Conversation):
+    def __init__(self, model_name, function_calling_llm : FunctionCallingLlm, category=None):
+        self.model_name = model_name
+        self.function_calling_llm = function_calling_llm
+        self.category = category
+    
+    def chat(self, message, role='user', print_output=True):
+        if role != 'user':
+            raise ValueError('Role must be "user"')
+
+        agent_id = os.environ.get('AGENT_ID', None)
+        
+        start_time = datetime.datetime.now()
+
+        response, usage_data = self.function_calling_llm.function_call_llm(message)
+
+        end_time = datetime.datetime.now()
+
+        print('Usage data:', usage_data)
+        if print_output:
+            print(response)
+        
+        #send_usage_to_db(usage_data, start_time, end_time, agent_id, self.category, self.model_name)
+        
+        return response
+
+class SambanovaToolformer(Toolformer):
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+
+    def new_conversation(self, system_prompt: str, tools: List[Tool], category=None) -> SambanovaConversation:
+        function_calling_llm = FunctionCallingLlm(system_prompt=system_prompt, tools=tools, select_expert=self.model_name)
+        return SambanovaConversation(self.model_name, function_calling_llm, category)
+
+#def make_llama_toolformer(model_name, system_prompt: str, tools: List[Tool]):
+#    if model_name not in ['llama3-8b', 'llama3-70b', 'llama3-405b']:
+#        raise ValueError(f"Unknown model name: {model_name}")
+#
+#    return SambanovaToolformer(model_name, system_prompt, tools)

toolformers/sambanova/function_calling.py

@@ -0,0 +1,312 @@
+import json
+import os
+import re
+from random import random
+from pprint import pprint
+import time
+from typing import List, Optional, Union
+
+from langchain_core.messages.ai import AIMessage
+from langchain_core.messages.human import HumanMessage
+from langchain_core.messages.tool import ToolMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableLambda
+
+from toolformers.base import Tool, StringParameter
+from toolformers.sambanova.api_gateway import APIGateway
+
+from toolformers.sambanova.utils import get_total_usage, usage_tracker
+
+
+FUNCTION_CALLING_SYSTEM_PROMPT = """You have access to the following tools:
+
+{tools}
+
+You can call one or more tools by adding a <ToolCalls> section to your message. For example:
+<ToolCalls>
+```json
+[{{
+  "tool": <name of the selected tool>,
+  "tool_input": <parameters for the selected tool, matching the tool's JSON schema>
+}}]
+```
+</ToolCalls>
+
+Note that you can select multiple tools at once by adding more objects to the list. Do not add \
+multiple <ToolCalls> sections to the same message.
+You will see the invocation of the tools in the response.
+
+
+Think step by step
+Do not call a tool if the input depends on another tool output that you do not have yet.
+Do not try to answer until you get all the tools output, if you do not have an answer yet, you can continue calling tools until you do.
+Your answer should be in the same language as the initial query.
+
+"""  # noqa E501
+
+
+conversational_response = Tool(
+    name='ConversationalResponse',
+    description='Respond conversationally only if no other tools should be called for a given query, or if you have a final answer. Response must be in the same language as the user query.',
+    parameters=[StringParameter(name='response', description='Conversational response to the user. Must be in the same language as the user query.', required=True)],
+    function=None
+)
+
+
+class FunctionCallingLlm:
+    """
+    function calling llm class
+    """
+
+    def __init__(
+        self,
+        tools: Optional[Union[Tool, List[Tool]]] = None,
+        default_tool: Optional[Tool] = None,
+        system_prompt: Optional[str] = None,
+        prod_mode: bool = False,
+        api: str = 'sncloud',
+        coe: bool = False,
+        do_sample: bool = False,
+        max_tokens_to_generate: Optional[int] = None,
+        temperature: float = 0.2,
+        select_expert: Optional[str] = None,
+    ) -> None:
+        """
+        Args:
+            tools (Optional[Union[Tool, List[Tool]]]): The tools to use.
+            default_tool (Optional[Tool]): The default tool to use.
+                defaults to ConversationalResponse
+            system_prompt (Optional[str]): The system prompt to use. defaults to FUNCTION_CALLING_SYSTEM_PROMPT
+            prod_mode (bool): Whether to use production mode. Defaults to False.
+            api (str): The api to use. Defaults to 'sncloud'.
+            coe (bool): Whether to use coe. Defaults to False.
+            do_sample (bool): Whether to do sample. Defaults to False.
+            max_tokens_to_generate (Optional[int]): The max tokens to generate. If None, the model will attempt to use the maximum available tokens.
+            temperature (float): The model temperature. Defaults to 0.2.
+            select_expert (Optional[str]): The expert to use. Defaults to None.
+        """
+        self.prod_mode = prod_mode
+        sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY')
+        self.api = api
+        self.llm = APIGateway.load_llm(
+            type=api,
+            streaming=True,
+            coe=coe,
+            do_sample=do_sample,
+            max_tokens_to_generate=max_tokens_to_generate,
+            temperature=temperature,
+            select_expert=select_expert,
+            process_prompt=False,
+            sambanova_api_key=sambanova_api_key,
+        )
+
+        if isinstance(tools, Tool):
+            tools = [tools]
+        self.tools = tools
+        if system_prompt is None:
+            system_prompt = ''
+
+        system_prompt = system_prompt.replace('{','{{').replace('}', '}}')
+        
+        if len(self.tools) > 0:
+            system_prompt += '\n\n'
+            system_prompt += FUNCTION_CALLING_SYSTEM_PROMPT
+        self.system_prompt = system_prompt
+
+        if default_tool is None:
+            default_tool = conversational_response
+
+    def execute(self, invoked_tools: List[dict]) -> tuple[bool, List[str]]:
+        """
+        Given a list of tool executions the llm return as required
+        execute them given the name with the mane in tools_map and the input arguments
+        if there is only one tool call and it is default conversational one, the response is marked as final response
+
+        Args:
+            invoked_tools (List[dict]): The list of tool executions generated by the LLM.
+        """
+        if self.tools is not None:
+            tools_map = {tool.name.lower(): tool for tool in self.tools}
+        else:
+            tools_map = {}
+        tool_msg = "Tool '{name}' response: {response}"
+        tools_msgs = []
+        if len(invoked_tools) == 1 and invoked_tools[0]['tool'].lower() == 'conversationalresponse':
+            final_answer = True
+            return final_answer, [invoked_tools[0]['tool_input']['response']]
+
+        final_answer = False
+
+        for tool in invoked_tools:
+            if tool['tool'].lower() == 'invocationerror':
+                tools_msgs.append(f'Tool invocation error: {tool["tool_input"]}')
+            elif tool['tool'].lower() != 'conversationalresponse':
+                print(f"\n\n---\nTool {tool['tool'].lower()} invoked with input {tool['tool_input']}\n")
+                
+                if tool['tool'].lower() not in tools_map:
+                    tools_msgs.append(f'Tool {tool["tool"]} not found')
+                else:
+                    response = tools_map[tool['tool'].lower()].call_tool_for_toolformer(**tool['tool_input'])
+                    # print(f'Tool response: {str(response)}\n---\n\n')
+                    tools_msgs.append(tool_msg.format(name=tool['tool'], response=str(response)))
+        return final_answer, tools_msgs
+
+    def json_finder(self, input_string: str) -> Optional[str]:
+        """
+        find json structures in an LLM string response, if bad formatted using LLM to correct it
+
+        Args:
+            input_string (str): The string to find the json structure in.
+        """
+
+        # 1. Ideal pattern: correctly surrounded by <ToolCalls> tags
+        json_pattern_1 = re.compile(r'<ToolCalls\>(.*)</ToolCalls\>', re.DOTALL + re.IGNORECASE)
+        # 2. Sometimes the closing tag is missing
+        json_pattern_2 = re.compile(r'<ToolCalls\>(.*)', re.DOTALL + re.IGNORECASE)
+        # 3. Sometimes it accidentally uses <ToolCall> instead of <ToolCalls>
+        json_pattern_3 = re.compile(r'<ToolCall\>(.*)</ToolCall\>', re.DOTALL + re.IGNORECASE)
+        # 4. Sometimes it accidentally uses <ToolCall> instead of <ToolCalls> and the closing tag is missing
+        json_pattern_4 = re.compile(r'<ToolCall\>(.*)', re.DOTALL + re.IGNORECASE)
+
+        # Find the first JSON structure in the string
+        json_match = json_pattern_1.search(input_string) or json_pattern_2.search(input_string) or json_pattern_3.search(input_string) or json_pattern_4.search(input_string)
+        if json_match:
+            json_str = json_match.group(1)
+
+            # 1. Outermost list of JSON object
+            call_pattern_1 = re.compile(r'\[.*\]', re.DOTALL)
+            # 2. Outermost JSON object
+            call_pattern_2 = re.compile(r'\{.*\}', re.DOTALL)
+
+            call_match_1 = call_pattern_1.search(json_str)
+            call_match_2 = call_pattern_2.search(json_str)
+
+            if call_match_1:
+                json_str = call_match_1.group(0)
+                try:
+                    return json.loads(json_str)
+                except Exception as e:
+                    return [{'tool': 'InvocationError', 'tool_input' : str(e)}]
+            elif call_match_2:
+                json_str = call_match_2.group(0)
+                try:
+                    return [json.loads(json_str)]
+                except Exception as e:
+                    return [{'tool': 'InvocationError', 'tool_input' : str(e)}]
+            else:
+                return [{'tool': 'InvocationError', 'tool_input' : 'Could not find JSON object in the <ToolCalls> section'}]
+        else:
+            dummy_json_response = [{'tool': 'ConversationalResponse', 'tool_input': {'response': input_string}}]
+            json_str = dummy_json_response
+        return json_str
+
+    def msgs_to_llama3_str(self, msgs: list) -> str:
+        """
+        convert a list of langchain messages with roles to expected LLmana 3 input
+
+        Args:
+            msgs (list): The list of langchain messages.
+        """
+        formatted_msgs = []
+        for msg in msgs:
+            if msg.type == 'system':
+                sys_placeholder = (
+                    '<|begin_of_text|><|start_header_id|>system<|end_header_id|>system<|end_header_id|> {msg}'
+                )
+                formatted_msgs.append(sys_placeholder.format(msg=msg.content))
+            elif msg.type == 'human':
+                human_placeholder = '<|eot_id|><|start_header_id|>user<|end_header_id|>\nUser: {msg} <|eot_id|><|start_header_id|>assistant<|end_header_id|>\nAssistant:'  # noqa E501
+                formatted_msgs.append(human_placeholder.format(msg=msg.content))
+            elif msg.type == 'ai':
+                assistant_placeholder = '<|eot_id|><|start_header_id|>assistant<|end_header_id|>\nAssistant: {msg}'
+                formatted_msgs.append(assistant_placeholder.format(msg=msg.content))
+            elif msg.type == 'tool':
+                tool_placeholder = '<|eot_id|><|start_header_id|>tools<|end_header_id|>\n{msg} <|eot_id|><|start_header_id|>assistant<|end_header_id|>\nAssistant:'  # noqa E501
+                formatted_msgs.append(tool_placeholder.format(msg=msg.content))
+            else:
+                raise ValueError(f'Invalid message type: {msg.type}')
+        return '\n'.join(formatted_msgs)
+
+    def msgs_to_sncloud(self, msgs: list) -> list:
+        """
+        convert a list of langchain messages with roles to expected FastCoE input
+
+        Args:
+            msgs (list): The list of langchain messages.
+        """
+        formatted_msgs = []
+        for msg in msgs:
+            if msg.type == 'system':
+                formatted_msgs.append({'role': 'system', 'content': msg.content})
+            elif msg.type == 'human':
+                formatted_msgs.append({'role': 'user', 'content': msg.content})
+            elif msg.type == 'ai':
+                formatted_msgs.append({'role': 'assistant', 'content': msg.content})
+            elif msg.type == 'tool':
+                formatted_msgs.append({'role': 'tools', 'content': msg.content})
+            else:
+                raise ValueError(f'Invalid message type: {msg.type}')
+        return json.dumps(formatted_msgs)
+
+    def function_call_llm(self, query: str, max_it: int = 10, debug: bool = False) -> str:
+        """
+        invocation method for function calling workflow
+
+        Args:
+            query (str): The query to execute.
+            max_it (int, optional): The maximum number of iterations. Defaults to 5.
+            debug (bool, optional): Whether to print debug information. Defaults to False.
+        """
+        function_calling_chat_template = ChatPromptTemplate.from_messages([('system', self.system_prompt)])
+        tools_schemas = [tool.as_llama_schema() for tool in self.tools]
+
+        history = function_calling_chat_template.format_prompt(tools=tools_schemas).to_messages()
+
+        history.append(HumanMessage(query))
+        tool_call_id = 0  # identification for each tool calling required to create ToolMessages
+        with usage_tracker():
+
+            for i in range(max_it):
+                json_parsing_chain = RunnableLambda(self.json_finder)
+
+                if self.api == 'sncloud':
+                    prompt = self.msgs_to_sncloud(history)
+                else:
+                    prompt = self.msgs_to_llama3_str(history)
+                # print(f'\n\n---\nCalling function calling LLM with prompt: \n{prompt}\n')
+                
+                exponential_backoff_lower = 30
+                exponential_backoff_higher = 60
+                llm_response = None
+                for _ in range(5):
+                    try:
+                        llm_response = self.llm.invoke(prompt, stream_options={'include_usage': True})
+                        break
+                    except Exception as e:
+                        if '429' in str(e):
+                            print('Rate limit exceeded. Waiting with random exponential backoff.')
+                            time.sleep(random() * (exponential_backoff_higher - exponential_backoff_lower) + exponential_backoff_lower)
+                            exponential_backoff_lower *= 2
+                            exponential_backoff_higher *= 2
+                        else:
+                            raise e
+
+                print('LLM response:', llm_response)
+
+                # print(f'\nFunction calling LLM response: \n{llm_response}\n---\n')
+                parsed_tools_llm_response = json_parsing_chain.invoke(llm_response)
+
+                history.append(AIMessage(llm_response))
+                final_answer, tools_msgs = self.execute(parsed_tools_llm_response)
+                if final_answer:  # if response was marked as final response in execution
+                    final_response = tools_msgs[0]
+                    if debug:
+                        print('\n\n---\nFinal function calling LLM history: \n')
+                        pprint(f'{history}')
+                    return final_response, get_total_usage()
+                else:
+                    history.append(ToolMessage('\n'.join(tools_msgs), tool_call_id=tool_call_id))
+                    tool_call_id += 1
+
+
+        raise Exception('Not a final response yet', history)

toolformers/sambanova/sambanova_langchain.py

@@ -0,0 +1,827 @@
+"""Langchain Wrapper around Sambanova LLM APIs."""
+
+import json
+from typing import Any, Dict, Generator, Iterator, List, Optional, Union
+
+import requests
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.language_models.llms import LLM
+from langchain_core.outputs import GenerationChunk
+from langchain_core.pydantic_v1 import Extra
+from langchain_core.utils import get_from_dict_or_env, pre_init
+from langchain_core.runnables import RunnableConfig, ensure_config
+from langchain_core.language_models.base import (
+    LanguageModelInput,
+)
+from toolformers.sambanova.utils import append_to_usage_tracker
+
+
+class SSEndpointHandler:
+    """
+    SambaNova Systems Interface for SambaStudio model endpoints.
+
+    :param str host_url: Base URL of the DaaS API service
+    """
+
+    def __init__(self, host_url: str, api_base_uri: str):
+        """
+        Initialize the SSEndpointHandler.
+
+        :param str host_url: Base URL of the DaaS API service
+        :param str api_base_uri: Base URI of the DaaS API service
+        """
+        self.host_url = host_url
+        self.api_base_uri = api_base_uri
+        self.http_session = requests.Session()
+
+    def _process_response(self, response: requests.Response) -> Dict:
+        """
+        Processes the API response and returns the resulting dict.
+
+        All resulting dicts, regardless of success or failure, will contain the
+        `status_code` key with the API response status code.
+
+        If the API returned an error, the resulting dict will contain the key
+        `detail` with the error message.
+
+        If the API call was successful, the resulting dict will contain the key
+        `data` with the response data.
+
+        :param requests.Response response: the response object to process
+        :return: the response dict
+        :type: dict
+        """
+        result: Dict[str, Any] = {}
+        try:
+            result = response.json()
+        except Exception as e:
+            result['detail'] = str(e)
+        if 'status_code' not in result:
+            result['status_code'] = response.status_code
+        return result
+
+    def _process_streaming_response(
+        self,
+        response: requests.Response,
+    ) -> Generator[Dict, None, None]:
+        """Process the streaming response"""
+        if 'api/predict/nlp' in self.api_base_uri:
+            try:
+                import sseclient
+            except ImportError:
+                raise ImportError(
+                    'could not import sseclient library' 'Please install it with `pip install sseclient-py`.'
+                )
+            client = sseclient.SSEClient(response)
+            close_conn = False
+            for event in client.events():
+                if event.event == 'error_event':
+                    close_conn = True
+                chunk = {
+                    'event': event.event,
+                    'data': event.data,
+                    'status_code': response.status_code,
+                }
+                yield chunk
+            if close_conn:
+                client.close()
+        elif 'api/v2/predict/generic' in self.api_base_uri or 'api/predict/generic' in self.api_base_uri:
+            try:
+                for line in response.iter_lines():
+                    chunk = json.loads(line)
+                    if 'status_code' not in chunk:
+                        chunk['status_code'] = response.status_code
+                    yield chunk
+            except Exception as e:
+                raise RuntimeError(f'Error processing streaming response: {e}')
+        else:
+            raise ValueError(f'handling of endpoint uri: {self.api_base_uri} not implemented')
+
+    def _get_full_url(self, path: str) -> str:
+        """
+        Return the full API URL for a given path.
+
+        :param str path: the sub-path
+        :returns: the full API URL for the sub-path
+        :type: str
+        """
+        return f'{self.host_url}/{self.api_base_uri}/{path}'
+
+    def nlp_predict(
+        self,
+        project: str,
+        endpoint: str,
+        key: str,
+        input: Union[List[str], str],
+        params: Optional[str] = '',
+        stream: bool = False,
+    ) -> Dict:
+        """
+        NLP predict using inline input string.
+
+        :param str project: Project ID in which the endpoint exists
+        :param str endpoint: Endpoint ID
+        :param str key: API Key
+        :param str input_str: Input string
+        :param str params: Input params string
+        :returns: Prediction results
+        :type: dict
+        """
+        if isinstance(input, str):
+            input = [input]
+        if 'api/predict/nlp' in self.api_base_uri:
+            if params:
+                data = {'inputs': input, 'params': json.loads(params)}
+            else:
+                data = {'inputs': input}
+        elif 'api/v2/predict/generic' in self.api_base_uri:
+            items = [{'id': f'item{i}', 'value': item} for i, item in enumerate(input)]
+            if params:
+                data = {'items': items, 'params': json.loads(params)}
+            else:
+                data = {'items': items}
+        elif 'api/predict/generic' in self.api_base_uri:
+            if params:
+                data = {'instances': input, 'params': json.loads(params)}
+            else:
+                data = {'instances': input}
+        else:
+            raise ValueError(f'handling of endpoint uri: {self.api_base_uri} not implemented')
+        response = self.http_session.post(
+            self._get_full_url(f'{project}/{endpoint}'),
+            headers={'key': key},
+            json=data,
+        )
+        return self._process_response(response)
+
+    def nlp_predict_stream(
+        self,
+        project: str,
+        endpoint: str,
+        key: str,
+        input: Union[List[str], str],
+        params: Optional[str] = '',
+    ) -> Iterator[Dict]:
+        """
+        NLP predict using inline input string.
+
+        :param str project: Project ID in which the endpoint exists
+        :param str endpoint: Endpoint ID
+        :param str key: API Key
+        :param str input_str: Input string
+        :param str params: Input params string
+        :returns: Prediction results
+        :type: dict
+        """
+        if 'api/predict/nlp' in self.api_base_uri:
+            if isinstance(input, str):
+                input = [input]
+            if params:
+                data = {'inputs': input, 'params': json.loads(params)}
+            else:
+                data = {'inputs': input}
+        elif 'api/v2/predict/generic' in self.api_base_uri:
+            if isinstance(input, str):
+                input = [input]
+            items = [{'id': f'item{i}', 'value': item} for i, item in enumerate(input)]
+            if params:
+                data = {'items': items, 'params': json.loads(params)}
+            else:
+                data = {'items': items}
+        elif 'api/predict/generic' in self.api_base_uri:
+            if isinstance(input, list):
+                input = input[0]
+            if params:
+                data = {'instance': input, 'params': json.loads(params)}
+            else:
+                data = {'instance': input}
+        else:
+            raise ValueError(f'handling of endpoint uri: {self.api_base_uri} not implemented')
+        # Streaming output
+        response = self.http_session.post(
+            self._get_full_url(f'stream/{project}/{endpoint}'),
+            headers={'key': key},
+            json=data,
+            stream=True,
+        )
+        for chunk in self._process_streaming_response(response):
+            yield chunk
+
+
+class SambaStudio(LLM):
+    """
+    SambaStudio large language models.
+
+    To use, you should have the environment variables
+    ``SAMBASTUDIO_BASE_URL`` set with your SambaStudio environment URL.
+    ``SAMBASTUDIO_BASE_URI`` set with your SambaStudio api base URI.
+    ``SAMBASTUDIO_PROJECT_ID`` set with your SambaStudio project ID.
+    ``SAMBASTUDIO_ENDPOINT_ID`` set with your SambaStudio endpoint ID.
+    ``SAMBASTUDIO_API_KEY``  set with your SambaStudio endpoint API key.
+
+    https://sambanova.ai/products/enterprise-ai-platform-sambanova-suite
+
+    read extra documentation in https://docs.sambanova.ai/sambastudio/latest/index.html
+
+    Example:
+    .. code-block:: python
+
+        from langchain_community.llms.sambanova  import SambaStudio
+        SambaStudio(
+            sambastudio_base_url="your-SambaStudio-environment-URL",
+            sambastudio_base_uri="your-SambaStudio-base-URI",
+            sambastudio_project_id="your-SambaStudio-project-ID",
+            sambastudio_endpoint_id="your-SambaStudio-endpoint-ID",
+            sambastudio_api_key="your-SambaStudio-endpoint-API-key,
+            streaming=False
+            model_kwargs={
+                "do_sample": False,
+                "max_tokens_to_generate": 1000,
+                "temperature": 0.7,
+                "top_p": 1.0,
+                "repetition_penalty": 1,
+                "top_k": 50,
+                #"process_prompt": False,
+                #"select_expert": "Meta-Llama-3-8B-Instruct"
+            },
+        )
+    """
+
+    sambastudio_base_url: str = ''
+    """Base url to use"""
+
+    sambastudio_base_uri: str = ''
+    """endpoint base uri"""
+
+    sambastudio_project_id: str = ''
+    """Project id on sambastudio for model"""
+
+    sambastudio_endpoint_id: str = ''
+    """endpoint id on sambastudio for model"""
+
+    sambastudio_api_key: str = ''
+    """sambastudio api key"""
+
+    model_kwargs: Optional[dict] = None
+    """Key word arguments to pass to the model."""
+
+    streaming: Optional[bool] = False
+    """Streaming flag to get streamed response."""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = 'forbid'#Extra.forbid
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        return True
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return {**{'model_kwargs': self.model_kwargs}}
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return 'Sambastudio LLM'
+
+    @pre_init
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values['sambastudio_base_url'] = get_from_dict_or_env(values, 'sambastudio_base_url', 'SAMBASTUDIO_BASE_URL')
+        values['sambastudio_base_uri'] = get_from_dict_or_env(
+            values,
+            'sambastudio_base_uri',
+            'SAMBASTUDIO_BASE_URI',
+            default='api/predict/generic',
+        )
+        values['sambastudio_project_id'] = get_from_dict_or_env(
+            values, 'sambastudio_project_id', 'SAMBASTUDIO_PROJECT_ID'
+        )
+        values['sambastudio_endpoint_id'] = get_from_dict_or_env(
+            values, 'sambastudio_endpoint_id', 'SAMBASTUDIO_ENDPOINT_ID'
+        )
+        values['sambastudio_api_key'] = get_from_dict_or_env(values, 'sambastudio_api_key', 'SAMBASTUDIO_API_KEY')
+        return values
+
+    def _get_tuning_params(self, stop: Optional[List[str]]) -> str:
+        """
+        Get the tuning parameters to use when calling the LLM.
+
+        Args:
+            stop: Stop words to use when generating. Model output is cut off at the
+                first occurrence of any of the stop substrings.
+
+        Returns:
+            The tuning parameters as a JSON string.
+        """
+        _model_kwargs = self.model_kwargs or {}
+        _kwarg_stop_sequences = _model_kwargs.get('stop_sequences', [])
+        _stop_sequences = stop or _kwarg_stop_sequences
+        # if not _kwarg_stop_sequences:
+        # _model_kwargs["stop_sequences"] = ",".join(
+        #    f'"{x}"' for x in _stop_sequences
+        # )
+        if 'api/v2/predict/generic' in self.sambastudio_base_uri:
+            tuning_params_dict = _model_kwargs
+        else:
+            tuning_params_dict = {k: {'type': type(v).__name__, 'value': str(v)} for k, v in (_model_kwargs.items())}
+        # _model_kwargs["stop_sequences"] = _kwarg_stop_sequences
+        tuning_params = json.dumps(tuning_params_dict)
+        return tuning_params
+
+    def _handle_nlp_predict(self, sdk: SSEndpointHandler, prompt: Union[List[str], str], tuning_params: str) -> str:
+        """
+        Perform an NLP prediction using the SambaStudio endpoint handler.
+
+        Args:
+            sdk: The SSEndpointHandler to use for the prediction.
+            prompt: The prompt to use for the prediction.
+            tuning_params: The tuning parameters to use for the prediction.
+
+        Returns:
+            The prediction result.
+
+        Raises:
+            ValueError: If the prediction fails.
+        """
+        response = sdk.nlp_predict(
+            self.sambastudio_project_id,
+            self.sambastudio_endpoint_id,
+            self.sambastudio_api_key,
+            prompt,
+            tuning_params,
+        )
+        if response['status_code'] != 200:
+            optional_detail = response.get('detail')
+            if optional_detail:
+                raise RuntimeError(
+                    f"Sambanova /complete call failed with status code "
+                    f"{response['status_code']}.\n Details: {optional_detail}"
+                )
+            else:
+                raise RuntimeError(
+                    f"Sambanova /complete call failed with status code "
+                    f"{response['status_code']}.\n response {response}"
+                )
+        if 'api/predict/nlp' in self.sambastudio_base_uri:
+            return response['data'][0]['completion']
+        elif 'api/v2/predict/generic' in self.sambastudio_base_uri:
+            return response['items'][0]['value']['completion']
+        elif 'api/predict/generic' in self.sambastudio_base_uri:
+            return response['predictions'][0]['completion']
+        else:
+            raise ValueError(f'handling of endpoint uri: {self.sambastudio_base_uri} not implemented')
+
+    def _handle_completion_requests(self, prompt: Union[List[str], str], stop: Optional[List[str]]) -> str:
+        """
+        Perform a prediction using the SambaStudio endpoint handler.
+
+        Args:
+            prompt: The prompt to use for the prediction.
+            stop: stop sequences.
+
+        Returns:
+            The prediction result.
+
+        Raises:
+            ValueError: If the prediction fails.
+        """
+        ss_endpoint = SSEndpointHandler(self.sambastudio_base_url, self.sambastudio_base_uri)
+        tuning_params = self._get_tuning_params(stop)
+        return self._handle_nlp_predict(ss_endpoint, prompt, tuning_params)
+
+    def _handle_nlp_predict_stream(
+        self, sdk: SSEndpointHandler, prompt: Union[List[str], str], tuning_params: str
+    ) -> Iterator[GenerationChunk]:
+        """
+        Perform a streaming request to the LLM.
+
+        Args:
+            sdk: The SVEndpointHandler to use for the prediction.
+            prompt: The prompt to use for the prediction.
+            tuning_params: The tuning parameters to use for the prediction.
+
+        Returns:
+            An iterator of GenerationChunks.
+        """
+        for chunk in sdk.nlp_predict_stream(
+            self.sambastudio_project_id,
+            self.sambastudio_endpoint_id,
+            self.sambastudio_api_key,
+            prompt,
+            tuning_params,
+        ):
+            if chunk['status_code'] != 200:
+                error = chunk.get('error')
+                if error:
+                    optional_code = error.get('code')
+                    optional_details = error.get('details')
+                    optional_message = error.get('message')
+                    raise ValueError(
+                        f"Sambanova /complete call failed with status code "
+                        f"{chunk['status_code']}.\n"
+                        f"Message: {optional_message}\n"
+                        f"Details: {optional_details}\n"
+                        f"Code: {optional_code}\n"
+                    )
+                else:
+                    raise RuntimeError(
+                        f"Sambanova /complete call failed with status code " f"{chunk['status_code']}." f"{chunk}."
+                    )
+            if 'api/predict/nlp' in self.sambastudio_base_uri:
+                text = json.loads(chunk['data'])['stream_token']
+            elif 'api/v2/predict/generic' in self.sambastudio_base_uri:
+                text = chunk['result']['items'][0]['value']['stream_token']
+            elif 'api/predict/generic' in self.sambastudio_base_uri:
+                if len(chunk['result']['responses']) > 0:
+                    text = chunk['result']['responses'][0]['stream_token']
+                else:
+                    text = ''
+            else:
+                raise ValueError(f'handling of endpoint uri: {self.sambastudio_base_uri}' f'not implemented')
+            generated_chunk = GenerationChunk(text=text)
+            yield generated_chunk
+
+    def _stream(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        """Call out to Sambanova's complete endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+        """
+        ss_endpoint = SSEndpointHandler(self.sambastudio_base_url, self.sambastudio_base_uri)
+        tuning_params = self._get_tuning_params(stop)
+        try:
+            if self.streaming:
+                for chunk in self._handle_nlp_predict_stream(ss_endpoint, prompt, tuning_params):
+                    if run_manager:
+                        run_manager.on_llm_new_token(chunk.text)
+                    yield chunk
+            else:
+                return
+        except Exception as e:
+            # Handle any errors raised by the inference endpoint
+            raise ValueError(f'Error raised by the inference endpoint: {e}') from e
+
+    def _handle_stream_request(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]],
+        run_manager: Optional[CallbackManagerForLLMRun],
+        kwargs: Dict[str, Any],
+    ) -> str:
+        """
+        Perform a streaming request to the LLM.
+
+        Args:
+            prompt: The prompt to generate from.
+            stop: Stop words to use when generating. Model output is cut off at the
+                first occurrence of any of the stop substrings.
+            run_manager: Callback manager for the run.
+            **kwargs: Additional keyword arguments. directly passed
+                to the sambastudio model in API call.
+
+        Returns:
+            The model output as a string.
+        """
+        completion = ''
+        for chunk in self._stream(prompt=prompt, stop=stop, run_manager=run_manager, **kwargs):
+            completion += chunk.text
+        return completion
+
+    def _call(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to Sambanova's complete endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+        """
+        if stop is not None:
+            raise Exception('stop not implemented')
+        try:
+            if self.streaming:
+                return self._handle_stream_request(prompt, stop, run_manager, kwargs)
+            return self._handle_completion_requests(prompt, stop)
+        except Exception as e:
+            # Handle any errors raised by the inference endpoint
+            raise ValueError(f'Error raised by the inference endpoint: {e}') from e
+
+
+class SambaNovaCloud(LLM):
+    """
+    SambaNova Cloud large language models.
+
+    To use, you should have the environment variables
+    ``SAMBANOVA_URL`` set with your SambaNova Cloud URL.
+    ``SAMBANOVA_API_KEY`` set with your SambaNova Cloud API Key.
+
+    http://cloud.sambanova.ai/
+
+    Example:
+    .. code-block:: python
+
+        SambaNovaCloud(
+            sambanova_url = SambaNova cloud endpoint URL,
+            sambanova_api_key = set with your SambaNova cloud API key,
+            max_tokens = mas number of tokens to generate
+            stop_tokens = list of stop tokens
+            model = model name
+        )
+    """
+
+    sambanova_url: str = ''
+    """SambaNova Cloud Url"""
+
+    sambanova_api_key: str = ''
+    """SambaNova Cloud api key"""
+
+    max_tokens: int = 4000
+    """max tokens to generate"""
+
+    stop_tokens: list = ['<|eot_id|>']
+    """Stop tokens"""
+
+    model: str = 'llama3-8b'
+    """LLM model expert to use"""
+
+    temperature: float = 0.0
+    """model temperature"""
+
+    top_p: float = 0.0
+    """model top p"""
+
+    top_k: int = 1
+    """model top k"""
+
+    stream_api: bool = True
+    """use stream api"""
+
+    stream_options: dict = {'include_usage': True}
+    """stream options, include usage to get generation metrics"""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = 'forbid'#Extra.forbid
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        return True
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            'model': self.model,
+            'max_tokens': self.max_tokens,
+            'stop': self.stop_tokens,
+            'temperature': self.temperature,
+            'top_p': self.top_p,
+            'top_k': self.top_k,
+        }
+
+    def invoke(
+        self,
+        input: LanguageModelInput,
+        config: Optional[RunnableConfig] = None,
+        *,
+        stop: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        config = ensure_config(config)
+
+        print('Invoking SambaNovaCloud with input:', input)
+        response = self.generate_prompt(
+                    [self._convert_input(input)],
+                    stop=stop,
+                    callbacks=config.get("callbacks"),
+                    tags=config.get("tags"),
+                    metadata=config.get("metadata"),
+                    run_name=config.get("run_name"),
+                    run_id=config.pop("run_id", None),
+                    **kwargs,
+                )
+        run_infos = response.run
+
+        if len(run_infos) > 1:
+            raise NotImplementedError('Multiple runs not supported')
+        
+        run_id = run_infos[0].run_id
+
+        #print('Raw response:', response.run)
+        #print('Run ID:', run_id)
+        #print(USAGE_TRACKER)
+        #if run_id in USAGE_TRACKER:
+        #    print('Usage:', USAGE_TRACKER[run_id])
+        #return response
+        return (
+            response
+            .generations[0][0]
+            .text
+        )
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return 'SambaNova Cloud'
+
+    @pre_init
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values['sambanova_url'] = get_from_dict_or_env(
+            values, 'sambanova_url', 'SAMBANOVA_URL', default='https://api.sambanova.ai/v1/chat/completions'
+        )
+        values['sambanova_api_key'] = get_from_dict_or_env(values, 'sambanova_api_key', 'SAMBANOVA_API_KEY')
+        return values
+
+    def _handle_nlp_predict_stream(
+        self,
+        prompt: Union[List[str], str],
+        stop: List[str],
+    ) -> Iterator[GenerationChunk]:
+        """
+        Perform a streaming request to the LLM.
+
+        Args:
+            prompt: The prompt to use for the prediction.
+            stop: list of stop tokens
+
+        Returns:
+            An iterator of GenerationChunks.
+        """
+        try:
+            import sseclient
+        except ImportError:
+            raise ImportError('could not import sseclient library' 'Please install it with `pip install sseclient-py`.')
+        try:
+            formatted_prompt = json.loads(prompt)
+        except:
+            formatted_prompt = [{'role': 'user', 'content': prompt}]
+
+        http_session = requests.Session()
+        if not stop:
+            stop = self.stop_tokens
+        data = {
+            'messages': formatted_prompt,
+            'max_tokens': self.max_tokens,
+            'stop': stop,
+            'model': self.model,
+            'temperature': self.temperature,
+            'top_p': self.top_p,
+            'top_k': self.top_k,
+            'stream': self.stream_api,
+            'stream_options': self.stream_options,
+        }
+        # Streaming output
+        response = http_session.post(
+            self.sambanova_url,
+            headers={'Authorization': f'Bearer {self.sambanova_api_key}', 'Content-Type': 'application/json'},
+            json=data,
+            stream=True,
+        )
+
+        client = sseclient.SSEClient(response)
+        close_conn = False
+
+        print('Response:', response)
+
+        if response.status_code != 200:
+            raise RuntimeError(
+                f'Sambanova /complete call failed with status code ' f'{response.status_code}.' f'{response.text}.'
+            )
+
+        for event in client.events():
+            if event.event == 'error_event':
+                close_conn = True
+            #print('Event:', event.data)
+            chunk = {
+                'event': event.event,
+                'data': event.data,
+                'status_code': response.status_code,
+            }
+
+            if chunk.get('error'):
+                raise RuntimeError(
+                    f"Sambanova /complete call failed with status code " f"{chunk['status_code']}." f"{chunk}."
+                )
+
+            try:
+                # check if the response is a final event in that case event data response is '[DONE]'
+                #if 'usage' in chunk['data']:
+                #    usage = json.loads(chunk['data'])
+                #    print('Usage:', usage)
+                if chunk['data'] != '[DONE]':
+                    data = json.loads(chunk['data'])
+                    if data.get('error'):
+                        raise RuntimeError(
+                            f"Sambanova /complete call failed with status code " f"{chunk['status_code']}." f"{chunk}."
+                        )
+                    # check if the response is a final response with usage stats (not includes content)
+                    if data.get('usage') is None:
+                        # check is not "end of text" response
+                        if data['choices'][0]['finish_reason'] is None:
+                            text = data['choices'][0]['delta']['content']
+                            generated_chunk = GenerationChunk(text=text)
+                            yield generated_chunk
+                    else:
+                        #if data['id'] not in USAGE_TRACKER:
+                        #    USAGE_TRACKER[data['id']] = []
+                        #USAGE_TRACKER[data['id']].append(data['usage'])
+                        append_to_usage_tracker(data['usage'])
+                        #print(f'Usage for id {data["id"]}:', data['usage'])
+            except Exception as e:
+                raise Exception(f'Error getting content chunk raw streamed response: {chunk}')
+
+    def _stream(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        """Call out to Sambanova's complete endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+        """
+        try:
+            for chunk in self._handle_nlp_predict_stream(prompt, stop):
+                if run_manager:
+                    run_manager.on_llm_new_token(chunk.text)
+                yield chunk
+        except Exception as e:
+            # Handle any errors raised by the inference endpoint
+            raise ValueError(f'Error raised by the inference endpoint: {e}') from e
+
+    def _handle_stream_request(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]],
+        run_manager: Optional[CallbackManagerForLLMRun],
+        kwargs: Dict[str, Any],
+    ) -> str:
+        """
+        Perform a streaming request to the LLM.
+
+        Args:
+            prompt: The prompt to generate from.
+            stop: Stop words to use when generating. Model output is cut off at the
+                first occurrence of any of the stop substrings.
+            run_manager: Callback manager for the run.
+            **kwargs: Additional keyword arguments. directly passed
+                to the Sambanova Cloud model in API call.
+
+        Returns:
+            The model output as a string.
+        """
+        completion = ''
+        for chunk in self._stream(prompt=prompt, stop=stop, run_manager=run_manager, **kwargs):
+            completion += chunk.text
+        return completion
+
+    def _call(
+        self,
+        prompt: Union[List[str], str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to Sambanova's  complete endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+        """
+        try:
+            return self._handle_stream_request(prompt, stop, run_manager, kwargs)
+        except Exception as e:
+            # Handle any errors raised by the inference endpoint
+            raise ValueError(f'Error raised by the inference endpoint: {e}') from e

toolformers/sambanova/utils.py

@@ -0,0 +1,32 @@
+from contextlib import contextmanager
+
+USAGE_TRACKER = None
+
+@contextmanager
+def usage_tracker():
+    global USAGE_TRACKER
+    assert USAGE_TRACKER is None
+    USAGE_TRACKER = []
+    try:
+        yield
+    finally:
+        USAGE_TRACKER = None
+
+def get_total_usage():
+    global USAGE_TRACKER
+
+    prompt_tokens = 0
+    completion_tokens = 0
+
+    for usage in USAGE_TRACKER:
+        prompt_tokens += usage['prompt_tokens']
+        completion_tokens += usage['completion_tokens']
+
+    return {
+        'prompt_tokens': prompt_tokens,
+        'completion_tokens': completion_tokens
+    }
+
+def append_to_usage_tracker(usage):
+    global USAGE_TRACKER
+    USAGE_TRACKER.append(usage)

utils.py

@@ -0,0 +1,80 @@
+import base64
+import hashlib
+
+from functools import wraps
+from typing import Callable, Any
+from inspect import signature, Parameter, Signature
+
+
+def extract_substring(text, start_tag, end_tag):
+    start_position = text.lower().find(start_tag.lower())
+    end_position = text.lower().find(end_tag.lower())
+
+    if start_position == -1 or end_position == -1:
+        return None
+
+    return text[start_position + len(start_tag):end_position].strip()
+
+def compute_hash(s):
+    # Hash a string using SHA-1 and return the base64 encoded result
+
+    m = hashlib.sha1()
+    m.update(s.encode())
+
+    b = m.digest()
+
+    return base64.b64encode(b).decode('ascii')
+
+def add_params_and_annotations(name: str, description: str, params: dict, return_type: Any):
+    """
+    A decorator to add parameters and a return type annotation to a function.
+
+    :param params: A dictionary where the keys are parameter names and values are their types.
+    :param return_type: The return type to add to the function's signature.
+    """
+    def decorator(func: Callable):
+        # Create new parameters based on the provided params dict
+        new_params = [
+            Parameter(name, Parameter.POSITIONAL_OR_KEYWORD, annotation=type_)
+            for name, (type_, _) in params.items()
+        ]
+        
+        # Get the existing signature and parameters
+        original_sig = signature(func)
+        original_params = list(original_sig.parameters.values())
+
+        # Combine new parameters with the existing ones
+        combined_params = new_params# + original_params
+
+        # Create a new signature with updated parameters and return annotation
+        new_sig = Signature(parameters=combined_params, return_annotation=return_type)
+        
+        # Define the wrapper function
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        docstring = description
+
+        if len(params) > 0:
+            docstring += '\n\nArgs:'
+            for param_name, (type_, param_description) in params.items():
+                docstring += f'\n  {param_name}: {param_description}'
+        
+        docstring += f'\n\nReturns:\n  {return_type.__name__}'
+
+        print('Input params:', params)
+
+
+        # Set the new signature on the wrapper
+        wrapper.__name__ = name
+        wrapper.__signature__ = new_sig
+        wrapper.__annotations__.update({ k: v[0] for k, v in params.items() })
+        wrapper.__annotations__['return'] = return_type
+        #wrapper.__annotations__['name'] = str
+        wrapper.__doc__ = docstring
+
+        print(docstring)
+
+        return wrapper
+    return decorator