File size: 1,015 Bytes
42a4238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | from typing import Any
from abc import ABC, abstractmethod
from pydantic import BaseModel, ConfigDict
from openai import OpenAI
import os
class Tool(ABC, BaseModel):
name: str
description: str
arg: str
def model_post_init(self, __context: Any) -> None:
self.name = self.name.lower().replace(' ', '_')
self.description = self.description.lower()
self.arg = self.arg.lower()
@abstractmethod
def run(self, prompt: str) -> str:
pass
def get_tool_description(self):
return f"Tool: {self.name}\nDescription: {self.description}\nArg: {self.arg}\n"
class LLMTool(Tool):
client: Any = None
def __init__(self, **data):
super().__init__(**data)
if self.client is None:
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY environment variable not set") # OPTIONAL : TAKE API-KEY AS INPUT AT THIS STAGE
self.client = OpenAI(api_key=api_key)
|