enable pass keyword args

Author: shahules786

nbs/llm/llm.ipynb

@@ -20,16 +20,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/opt/homebrew/Caskroom/miniforge/base/envs/random/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "#| export\n",
     "\n",
@@ -42,9 +33,10 @@
     "T = t.TypeVar('T', bound=BaseModel)\n",
     "\n",
     "class RagasLLM:\n",
-    "    def __init__(self, provider: str, model:str, client: t.Any):\n",
+    "    def __init__(self, provider: str, model:str, client: t.Any, **model_args):\n",
     "        self.provider = provider.lower()\n",
     "        self.model = model\n",
+    "        self.model_args = model_args or {}\n",
     "        self.client = self._initialize_client(provider, client)\n",
     "        # Check if client is async-capable at initialization\n",
     "        self.is_async = self._check_client_async()\n",
@@ -98,33 +90,30 @@
     "                loop.close()\n",
     "                asyncio.set_event_loop(None)\n",
     "    \n",
-    "    def generate(self, prompt: str, response_model: t.Type[T], **kwargs) -> T:\n",
+    "    def generate(self, prompt: str, response_model: t.Type[T]) -> T:\n",
     "        \"\"\"Generate a response using the configured LLM.\n",
     "        \n",
     "        For async clients, this will run the async method in the appropriate event loop.\n",
     "        \"\"\"\n",
     "        messages = [{\"role\": \"user\", \"content\": prompt}]\n",
-    "        if \"model\" not in kwargs and self.model:\n",
-    "            kwargs[\"model\"] = self.model\n",
     "        \n",
     "        # If client is async, use the appropriate method to run it\n",
     "        if self.is_async:\n",
     "            return self._run_async_in_current_loop(\n",
-    "                self.agenerate(prompt, response_model, **kwargs)\n",
+    "                self.agenerate(prompt, response_model)\n",
     "            )\n",
     "        else:\n",
     "            # Regular sync client, just call the method directly\n",
     "            return self.client.chat.completions.create(\n",
+    "                model=self.model,\n",
     "                messages=messages,\n",
     "                response_model=response_model,\n",
-    "                **kwargs\n",
+    "                **self.model_args,\n",
     "            )\n",
     "    \n",
-    "    async def agenerate(self, prompt: str, response_model: t.Type[T], **kwargs) -> T:\n",
+    "    async def agenerate(self, prompt: str, response_model: t.Type[T]) -> T:\n",
     "        \"\"\"Asynchronously generate a response using the configured LLM.\"\"\"\n",
     "        messages = [{\"role\": \"user\", \"content\": prompt}]\n",
-    "        if \"model\" not in kwargs and self.model:\n",
-    "            kwargs[\"model\"] = self.model\n",
     "        \n",
     "        # If client is not async, raise a helpful error\n",
     "        if not self.is_async:\n",
@@ -134,13 +123,14 @@
     "        \n",
     "        # Regular async client, call the method directly\n",
     "        return await self.client.chat.completions.create(\n",
+    "            model=self.model,\n",
     "            messages=messages,\n",
     "            response_model=response_model,\n",
-    "            **kwargs\n",
+    "            **self.model_args,\n",
     "        )\n",
     "\n",
-    "def ragas_llm(provider: str,model:str, client: t.Any,) -> RagasLLM:\n",
-    "    return RagasLLM(provider=provider, client=client, model=model)"
+    "def ragas_llm(provider: str,model:str, client: t.Any, **model_args) -> RagasLLM:\n",
+    "    return RagasLLM(provider=provider, client=client, model=model, **model_args)"
    ]
   },
   {
@@ -201,8 +191,26 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Response(response='The capital of India is New Delhi.')"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#| eval: false\n",
+    "\n",
+    "from anthropic import Anthropic\n",
+    "\n",
+    "llm = ragas_llm(provider=\"anthropic\",model=\"claude-3-opus-20240229\",client=Anthropic(),max_tokens=1024)\n",
+    "llm.generate(\"What is the capital of India?\",response_model=Response)"
+   ]
   }
  ],
  "metadata": {

ragas_annotator/llm/llm.py

@@ -13,9 +13,10 @@
 T = t.TypeVar('T', bound=BaseModel)
 
 class RagasLLM:
-    def __init__(self, provider: str, model:str, client: t.Any):
+    def __init__(self, provider: str, model:str, client: t.Any, **model_args):
         self.provider = provider.lower()
         self.model = model
+        self.model_args = model_args or {}
         self.client = self._initialize_client(provider, client)
         # Check if client is async-capable at initialization
         self.is_async = self._check_client_async()
@@ -69,33 +70,30 @@ def _run_async_in_current_loop(self, coro):
                 loop.close()
                 asyncio.set_event_loop(None)
 
-    def generate(self, prompt: str, response_model: t.Type[T], **kwargs) -> T:
+    def generate(self, prompt: str, response_model: t.Type[T]) -> T:
         """Generate a response using the configured LLM.
         
         For async clients, this will run the async method in the appropriate event loop.
         """
         messages = [{"role": "user", "content": prompt}]
-        if "model" not in kwargs and self.model:
-            kwargs["model"] = self.model
 
         # If client is async, use the appropriate method to run it
         if self.is_async:
             return self._run_async_in_current_loop(
-                self.agenerate(prompt, response_model, **kwargs)
+                self.agenerate(prompt, response_model)
             )
         else:
             # Regular sync client, just call the method directly
             return self.client.chat.completions.create(
+                model=self.model,
                 messages=messages,
                 response_model=response_model,
-                **kwargs
+                **self.model_args,
             )
 
-    async def agenerate(self, prompt: str, response_model: t.Type[T], **kwargs) -> T:
+    async def agenerate(self, prompt: str, response_model: t.Type[T]) -> T:
         """Asynchronously generate a response using the configured LLM."""
         messages = [{"role": "user", "content": prompt}]
-        if "model" not in kwargs and self.model:
-            kwargs["model"] = self.model
 
         # If client is not async, raise a helpful error
         if not self.is_async:
@@ -105,10 +103,11 @@ async def agenerate(self, prompt: str, response_model: t.Type[T], **kwargs) -> T
 
         # Regular async client, call the method directly
         return await self.client.chat.completions.create(
+            model=self.model,
             messages=messages,
             response_model=response_model,
-            **kwargs
+            **self.model_args,
         )
 
-def ragas_llm(provider: str,model:str, client: t.Any,) -> RagasLLM:
-    return RagasLLM(provider=provider, client=client, model=model)
+def ragas_llm(provider: str,model:str, client: t.Any, **model_args) -> RagasLLM:
+    return RagasLLM(provider=provider, client=client, model=model, **model_args)