From a9d9e13474b29b05e9f2c537b3dd34b091bb33c5 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 1 Sep 2024 11:59:43 +0200 Subject: [PATCH] Adapter class added --- README.md | 3 +- data_folder/config.yaml | 8 ++--- src/gpt.py | 79 +++++++++++++++++++++++++++++++++-------- 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index b396b3f..fc4f060 100644 --- a/README.md +++ b/README.md @@ -211,11 +211,12 @@ This file defines your job search parameters and bot behavior. Each section cont - Marketing ``` - `llm_model_type`: - - Choose the model type, supported: openai / ollama + - Choose the model type, supported: openai / ollama / claude - `llm_model`: - Choose the LLM model, currently supported: - openai: gpt-4o - ollama: llama2, mistral:v0.3 + - claude: any model - `llm_api_url`: - Link of the API endpoint for the LLM model diff --git a/data_folder/config.yaml b/data_folder/config.yaml index e28e58a..60647c0 100644 --- a/data_folder/config.yaml +++ b/data_folder/config.yaml @@ -40,7 +40,7 @@ companyBlacklist: titleBlacklist: - word1 - word2 - -llm_model_type: [openai / ollama] -llm_model: ['gpt-4o' / 'mistral:v0.3'] -llm_api_url: [https://api.pawan.krd/cosmosrp/v1', http://127.0.0.1:11434/] \ No newline at end of file + +llm_model_type: [openai / ollama / claude] +llm_model: ['gpt-4o' / 'mistral:v0.3' / anymodel] +llm_api_url: [https://api.pawan.krd/cosmosrp/v1' / http://127.0.0.1:11434/] \ No newline at end of file diff --git a/src/gpt.py b/src/gpt.py index bc71e8b..b4b66bc 100644 --- a/src/gpt.py +++ b/src/gpt.py @@ -3,6 +3,7 @@ import os import re import textwrap from datetime import datetime +from abc import ABC, abstractmethod from typing import Dict, List, Union from pathlib import Path from dotenv import load_dotenv @@ -11,17 +12,75 @@ from langchain_core.output_parsers import StrOutputParser from langchain_core.prompt_values import StringPromptValue from langchain_core.prompts import ChatPromptTemplate from langchain_openai import ChatOpenAI -from langchain_ollama import ChatOllama from Levenshtein import distance import src.strings as strings load_dotenv() +class AIModel(ABC): + @abstractmethod + def generate_response(self, prompt: str) -> str: + pass + +class OpenAIModel(AIModel): + def __init__(self, api_key: str, llm_model: str, llm_api_url: str): + from langchain_openai import ChatOpenAI + self.model = ChatOpenAI(model_name=llm_model, openai_api_key=api_key, + temperature=0.4, base_url=llm_api_url) + + def generate_response(self, prompt: str) -> str: + response = self.model.invoke(prompt) + return response.content + +class ClaudeModel(AIModel): + def __init__(self, api_key: str, llm_model: str, llm_api_url: str): + from anthropic import Anthropic + self.client = Anthropic(api_key=api_key) + + def generate_response(self, prompt: str) -> str: + formatted_prompt = f"\n\nHuman: {prompt}\n\nAssistant:" + response = self.client.completions.create( + model="claude-2", + prompt=formatted_prompt, + max_tokens_to_sample=300 + ) + return response.completion.strip() + +class OllamaModel(AIModel): + def __init__(self, api_key: str, llm_model: str, llm_api_url: str): + from langchain_ollama import ChatOllama + self.model = ChatOllama(model=llm_model, base_url=llm_api_url) + + def generate_response(self, prompt: str) -> str: + response = self.model.invoke(prompt) + return response.content + +class AIAdapter: + def __init__(self, config: dict, api_key: str): + self.model = self._create_model(config, api_key) + + def _create_model(self, config: dict, api_key: str) -> AIModel: + llm_model_type = config['llm_model_type'] + llm_model = config['llm_model'] + llm_api_url = config['llm_api_url'] + print('Using {0} with {1} from {2}'.format(llm_model_type, llm_model, llm_api_url)) + + if llm_model_type == "openai": + return OpenAIModel(api_key, llm_model, llm_api_url) + elif llm_model_type == "claude": + return ClaudeModel(api_key, llm_model, llm_api_url) + elif llm_model_type == "ollama": + return OllamaModel(api_key, llm_model, llm_api_url) + else: + raise ValueError(f"Unsupported model type: {model_type}") + + def generate_response(self, prompt: str) -> str: + return self.model.generate_response(prompt) class LLMLogger: - def __init__(self, llm: Union[ChatOpenAI, ChatOllama]): + def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel]): self.llm = llm @staticmethod @@ -79,7 +138,7 @@ class LLMLogger: class LoggerChatModel: - def __init__(self, llm: Union[ChatOpenAI, ChatOllama]): + def __init__(self, llm: Union[OpenAIModel, OllamaModel, ClaudeModel]): self.llm = llm def __call__(self, messages: List[Dict[str, str]]) -> str: @@ -115,18 +174,8 @@ class LoggerChatModel: class GPTAnswerer: def __init__(self, config, llm_api_key): - llm_model_type = config['llm_model_type'] - llm_model = config['llm_model'] - llm_api_url = config['llm_api_url'] - - print('Using {0} with {1} from {2}'.format(llm_model_type, llm_model, llm_api_url)) - - if llm_model_type == "ollama": - self.llm_model = ChatOllama(model=llm_model, temperature = 0.4, base_url=llm_api_url) - elif llm_model_type == "openai": - self.llm_model = ChatOpenAI(model_name=llm_model, openai_api_key=llm_api_key, temperature=0.4, - base_url=llm_api_url) - self.llm_cheap = LoggerChatModel(self.llm_model) + self.ai_adapter = AIAdapter(config, llm_api_key) + self.llm_cheap = LoggerChatModel(self.ai_adapter) @property def job_description(self): return self.job.description