Initial commit

pkg/llm/__init__.py

@@ -0,0 +1,28 @@
+import aiohttp
+
+from .data import APIEndpoints, Payload
+from .exception import LanguageModelException
+from .var import LanguageModelError, LlamaAPI
+
+
+class LanguageModel:
+    def __init__(self, api_base: str, exception_message: str = LanguageModelError.GENERIC_EXCEPTION):
+        self.endpoints = APIEndpoints(
+            api_base,
+            LlamaAPI.CHAT_COMPLETIONS
+        )
+        self.exception_message = exception_message
+        
+    async def make_request(self, payload: Payload) -> str:
+        data = payload.render()
+        
+        async with aiohttp.ClientSession() as session:
+            try:
+                async with session.post(self.endpoints.chat_completion, json = data) as response:
+                    if response.status == 200:
+                        result = await response.json()
+                        return result["choices"][0]["message"]["content"]
+                    raise LanguageModelException(f"Request failed with status code {response.status}!")
+            except (aiohttp.ClientError, LanguageModelException) as e:
+                print(f"Exception occurred: {e}")
+                return self.exception_message

pkg/llm/data.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+
+from .var import Chat, LLMParam, Who
+
+
+@dataclass()
+class Payload:
+    system_prompt: str
+    user_message: str
+    temperature: float
+    max_tokens: int
+    stream: bool
+    
+    def render(self) -> dict:
+        return {
+            LLMParam.MESSAGES: [
+                {Chat.ROLE: Who.SYSTEM, Chat.CONTENT: self.system_prompt},
+                {Chat.ROLE: Who.USER, Chat.CONTENT: self.user_message}
+            ],
+            LLMParam.TEMPERATURE: self.temperature,
+            LLMParam.MAX_TOKENS: self.max_tokens,
+            LLMParam.STREAM: self.stream
+        }
+        
+@dataclass
+class APIEndpoints:
+    api_base: str
+    chat_completion: str
+    
+    def __post_init__(self):
+        self.api_base = self.api_base.rstrip("/")
+        
+        for field_name in self.__dataclass_fields__:
+            if field_name == "api_base":
+                continue
+                
+            path = getattr(self, field_name)
+            clean_path = f"/{path.lstrip('/')}"
+            
+            setattr(self, field_name, f"{self.api_base}{clean_path}")

pkg/llm/exception.py

@@ -0,0 +1,2 @@
+class LanguageModelException(Exception):
+    pass

pkg/llm/var.py

@@ -0,0 +1,23 @@
+from enum import StrEnum
+
+
+class LLMParam(StrEnum):
+    MESSAGES = "messages"
+    TEMPERATURE = "temperature"
+    MAX_TOKENS = "max_tokens"
+    STREAM = "stream"
+    
+class Chat(StrEnum):
+    ROLE = "role"
+    CONTENT = "content"
+    
+class Who(StrEnum):
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    
+class LlamaAPI(StrEnum):
+    CHAT_COMPLETIONS = "/v1/chat/completions"
+
+class LanguageModelError(StrEnum):
+    GENERIC_EXCEPTION = "Language model server could not be contacted!"