update

llm.py

@@ -1,6 +1,7 @@
 import urllib3
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-import os
+
+from pydantic import Extra
 import requests
 from typing import Any, List, Mapping, Optional
 
@@ -9,8 +10,9 @@ from langchain.llms.base import LLM
 
 class URALLM(LLM):
     llm_url = os.environ.get("URL")
+
     class Config:
-        extra = 'forbid'
+        extra = Extra.forbid
 
     @property
     def _llm_type(self) -> str:
@@ -35,11 +37,10 @@ class URALLM(LLM):
                            "repetition_penalty":1.1,
                            "do_sample":True,
                            "top_k":10
-                           },
-            "token": os.environ.get("TOKEN")
+                           }
         }
 
-        headers = {"Content-Type": "application/json"}
+        headers = {"Content-Type": "application/json", "Authorization": os.environ.get("TOKEN")}
 
         response = requests.post(self.llm_url, json=payload, headers=headers, verify=False)
         response.raise_for_status()
@@ -52,4 +53,5 @@ class URALLM(LLM):
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
         """Get the identifying parameters."""
-        return {"llmUrl": self.llm_url}
+        return {"llmUrl": self.llm_url}
+

rag.py

@@ -110,9 +110,9 @@ class CustomRetriever(BaseRetriever):
         for idx in sim_scores_argsort:
           docs.append(documents[idx])
 
-        docs_top_4 = docs[0:4]
+        docs_top_2 = docs[0:2]
 
-        return docs_top_4
+        return docs_top_2
     
 llm = URALLM()
 custom_retriever = CustomRetriever(vectorstores = vectordb,retriever = vectordb.as_retriever(search_kwargs={"k": 50}))