Skip to content

Instantly share code, notes, and snippets.

@lamoboos223
Last active December 16, 2024 22:11
Show Gist options
  • Save lamoboos223/df5b9318d456af0ee82261acb7b1ebfb to your computer and use it in GitHub Desktop.
Save lamoboos223/df5b9318d456af0ee82261acb7b1ebfb to your computer and use it in GitHub Desktop.
This is a demo on how to make a function calling app via LangChain framework to utilize models running on Ollama platform
import warnings
from fastapi import FastAPI
from langchain_experimental.llms.ollama_functions import OllamaFunctions
from pydantic import BaseModel
from typing import List

warnings.filterwarnings('ignore')

app = FastAPI()

# Initialize the model globally
model = OllamaFunctions(
    model="llama3", 
    format="json"
)

model = model.bind_tools(
    tools=[
        {
            "name": "get_current_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                    },
                },
                "required": ["location"],
            },
        },
        {
            "name": "check_stock_market", 
            "description": "Get the current stock price and market information",
            "parameters": {
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "The stock symbol/ticker, e.g. GOOGL",
                    },
                    "info_type": {
                        "type": "string",
                        "enum": ["price", "volume", "market_cap", "full"],
                        "description": "Type of stock information to retrieve"
                    }
                },
                "required": ["symbol"],
            },
        }
    ]
)

class Query(BaseModel):
    question: str

class ToolCall(BaseModel):
    name: str
    args: dict
    id: str
    type: str

class Response(BaseModel):
    tool_calls: List[ToolCall]
    details: dict

@app.post("/ask", response_model=Response)
async def ask_question(query: Query):
    response = model.invoke(query.question)
    
    tool_names = []
    tool_args = []
    tool_ids = []
    tool_types = []
    
    for tool_call in response.tool_calls:
        tool_names.append(tool_call['name'])
        tool_args.append(tool_call['args'])
        tool_ids.append(tool_call['id'])
        tool_types.append(tool_call['type'])

    details = {}
    if tool_names[0] == "get_current_weather":
        details = {
            "location": tool_args[0]['location'],
            "unit": tool_args[0].get('unit', 'celsius')  # Default to celsius if not specified
        }
    elif tool_names[0] == "check_stock_market":
        details = {
            "symbol": tool_args[0]['symbol'],
            "info_type": tool_args[0].get('info_type', 'price')  # Default to price if not specified
        }

    return Response(
        tool_calls=[
            ToolCall(
                name=name,
                args=args,
                id=id,
                type=type
            ) for name, args, id, type in zip(tool_names, tool_args, tool_ids, tool_types)
        ],
        details=details
    )
@lamoboos223
Copy link
Author

Test

image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment