From f89398aaf65dff90d627e2952866d3d30bdde1fb Mon Sep 17 00:00:00 2001
From: Tarun Jain <132284203+tarun-aiplanet@users.noreply.github.com>
Date: Fri, 21 Jun 2024 01:37:29 +0530
Subject: [PATCH] claude sonnet integration (#59)

Co-authored-by: lucifertrj <tarun.19ec144@cambridge.edu.in>
---
 src/beyondllm/llms/__init__.py |  3 +-
 src/beyondllm/llms/claude.py   | 85 ++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 src/beyondllm/llms/claude.py

diff --git a/src/beyondllm/llms/__init__.py b/src/beyondllm/llms/__init__.py
index 5080768..44bd1d7 100644
--- a/src/beyondllm/llms/__init__.py
+++ b/src/beyondllm/llms/__init__.py
@@ -6,4 +6,5 @@
 from .ollama import OllamaModel
 from .multimodal import GeminiMultiModal
 from .gpt4o import GPT4oOpenAIModel
-from .chatgroq import GroqModel
\ No newline at end of file
+from .chatgroq import GroqModel
+from .claude import ClaudeModel
\ No newline at end of file
diff --git a/src/beyondllm/llms/claude.py b/src/beyondllm/llms/claude.py
new file mode 100644
index 0000000..72cca59
--- /dev/null
+++ b/src/beyondllm/llms/claude.py
@@ -0,0 +1,85 @@
+from beyondllm.llms.base import BaseLLMModel, ModelConfig
+from typing import Any, Dict, List, Optional
+import os
+from dataclasses import dataclass,field
+
+@dataclass
+class ClaudeModel:
+    """
+    Class representing a Claude model from Antrophic
+    Example:
+    from beyondllm.llms import ClaudeModel
+    llm = ClaudeModel(model="claude-3-5-sonnet-20240620",api_key = "",model_kwargs = {"max_tokens":512,"temperature":0.1})
+    """
+    api_key: str = ""
+    model: str = field(default="claude-3-5-sonnet-20240620")
+    model_kwargs: dict = field(default_factory=lambda: {
+                    "temperature": 0,
+                    "top_p": 1,
+                    "top_k": 1,
+                    "max_tokens": 1000,
+            })
+
+    def __post_init__(self):
+        if not self.api_key:  
+            self.api_key = os.getenv('ANTHROPIC_API_KEY') 
+            if not self.api_key: 
+                raise ValueError("ANTHROPIC_API_KEY is not provided and not found in environment variables.")
+        self.load_llm()
+
+    def load_llm(self):
+        try:
+            import anthropic
+        except ImportError:
+            raise ImportError("Anthropic library is not installed. Please install it with ``pip install anthropic``.")
+        
+        try:
+            self.client = anthropic.Anthropic(api_key=self.api_key)
+
+        except Exception as e:
+            raise Exception("Failed to load the model from Anthropic:", str(e))
+        
+        return self.client
+
+    def predict(self,prompt:Any):
+        if self.model_kwargs is not None:
+            response = self.client.messages.create(
+                model = self.model,
+                **self.model_kwargs,
+                system="You are a world-class AI Assistant",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": prompt
+                            }
+                        ]
+                    }
+                ]
+            )
+        else:
+            response = self.client.messages.create(
+                model = self.model,
+                system="You are a world-class AI Assistant",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": prompt
+                            }
+                        ]
+                    }
+                ]
+            )
+        
+        return response.content[0].text
+    
+    @staticmethod
+    def load_from_kwargs(self,kwargs): 
+        model_config = ModelConfig(**kwargs)
+        self.config = model_config
+        self.load_llm()
\ No newline at end of file