'Refactored by Sourcery'

libs/legacy/app/api/agent_documents.py

@@ -17,15 +17,11 @@
 
 def parse_filter_params(request: Request):
     query_params = request.query_params
-    filter_params = {}
-
-    for k, v in query_params.items():
-        if k.startswith("filter[") and k.endswith("]"):
-            # Removing 'filter[' from start and ']' from end
-            filter_key = k[7:-1]
-            filter_params[filter_key] = v
-
-    return filter_params
+    return {
+        k[7:-1]: v
+        for k, v in query_params.items()
+        if k.startswith("filter[") and k.endswith("]")
+    }
 
 
 @router.post(

libs/legacy/app/api/agent_tools.py

@@ -14,15 +14,11 @@
 
 def parse_filter_params(request: Request):
     query_params = request.query_params
-    filter_params = {}
-
-    for k, v in query_params.items():
-        if k.startswith("filter[") and k.endswith("]"):
-            # Removing 'filter[' from start and ']' from end
-            filter_key = k[7:-1]
-            filter_params[filter_key] = v
-
-    return filter_params
+    return {
+        k[7:-1]: v
+        for k, v in query_params.items()
+        if k.startswith("filter[") and k.endswith("]")
+    }
 
 
 @router.post(

libs/legacy/app/api/agents.py

@@ -119,9 +119,9 @@ async def read_library_agents(token=Depends(JWTBearer())):
 )
 async def read_agent(agentId: str, token=Depends(JWTBearer())):
     """Agent detail endpoint"""
-    agent = prisma.agent.find_unique(where={"id": agentId}, include={"prompt": True})
-
-    if agent:
+    if agent := prisma.agent.find_unique(
+        where={"id": agentId}, include={"prompt": True}
+    ):
         return {"success": True, "data": agent}
 
     logger.error(f"Agent with id: {agentId} not found")

libs/legacy/app/api/auth.py

@@ -19,13 +19,12 @@
 @router.post("/auth/sign-in", response_model=SignInOutput)
 async def sign_in(signIn: SignIn):
     try:
-        user = prisma.user.find_first(
+        if user := prisma.user.find_first(
             where={
                 "email": signIn.email,
             },
             include={"profile": True},
-        )
-        if user:
+        ):
             validated = validatePassword(signIn.password, user.password)
             del user.password
 
@@ -34,16 +33,12 @@ async def sign_in(signIn: SignIn):
                 return {"success": True, "data": SignInOut(token=token, user=user)}
 
             logger.warning("Invalid password")
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Invalid credentials",
-            )
         else:
             logger.warning("User not found")
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Invalid credentials",
-            )
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid credentials",
+        )
     except Exception as e:
         logger.error("Couldn't find user by email", exc_info=e)
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
@@ -96,16 +91,15 @@ async def oauth_handler(body: OAuth):
     )
     if user:
         return {"success": True, "data": user}
-    else:
-        user = prisma.user.create(
-            {
-                "email": body.email,
-                "provider": body.provider,
-                "name": body.name,
-                "accessToken": body.access_token,
-            }
-        )
-        prisma.profile.create({"userId": user.id})
+    user = prisma.user.create(
+        {
+            "email": body.email,
+            "provider": body.provider,
+            "name": body.name,
+            "accessToken": body.access_token,
+        }
+    )
+    prisma.profile.create({"userId": user.id})
 
-        if user:
-            return {"success": True, "data": user}
+    if user:
+        return {"success": True, "data": user}

libs/legacy/app/api/users.py

@@ -34,8 +34,9 @@ async def read_user_me(token=Depends(JWTBearer())):
 @router.get("/users/{userId}", response_model=UserOutput)
 async def read_user(userId: str):
     try:
-        user = prisma.user.find_unique(where={"id": userId}, include={"profile": True})
-        if user:
+        if user := prisma.user.find_unique(
+            where={"id": userId}, include={"profile": True}
+        ):
             return {"success": True, "data": user}
         else:
             logger.error("Couldn't find user")

libs/legacy/app/lib/agents/agent.py

@@ -20,16 +20,14 @@ def get_agent(self, session: str | None = None) -> AgentExecutor:
         llm = self.agent_base._get_llm()
         memory = self.agent_base._get_memory(session)
         prompt = self.agent_base._get_prompt()
-        agent = LLMChain(
+        return LLMChain(
             llm=llm,
             memory=memory,
             verbose=True,
             prompt=prompt,
             output_key="output",
         )
 
-        return agent
-
 
 class OpenAIAgent(AgentStrategy):
     def __init__(self, agent_base: AgentBase):
@@ -40,7 +38,7 @@ def get_agent(self, session: str | None = None) -> AgentExecutor:
         tools = self.agent_base._get_tools()
         memory = self.agent_base._get_memory(session)
         prompt = self.agent_base._get_prompt()
-        agent = initialize_agent(
+        return initialize_agent(
             tools=tools,
             llm=llm,
             agent=AgentType.OPENAI_FUNCTIONS,
@@ -58,8 +56,6 @@ def get_agent(self, session: str | None = None) -> AgentExecutor:
             return_intermediate_steps=True,
         )
 
-        return agent
-
 
 class ReactAgent(AgentStrategy):
     def __init__(self, agent_base):
@@ -79,12 +75,10 @@ def get_agent(self, session: str | None = None) -> AgentExecutor:
             stop=["\nObservation:"],
             allowed_tools=tool_names,
         )
-        agent = AgentExecutor.from_agent_and_tools(
+        return AgentExecutor.from_agent_and_tools(
             agent=agent_config,
             tools=tools,
             verbose=True,
             memory=memory,
             return_intermediate_steps=True,
         )
-
-        return agent

libs/legacy/app/lib/agents/base.py

@@ -86,7 +86,7 @@ def __init__(
         self.tools = self._get_agent_tools()
 
     def _get_api_key(self) -> str:
-        if self.llm["provider"] == "openai-chat" or self.llm["provider"] == "openai":
+        if self.llm["provider"] in ["openai-chat", "openai"]:
             return (
                 self.llm["api_key"]
                 if "api_key" in self.llm
@@ -161,7 +161,7 @@ def _get_llm(self, has_streaming: bool = True) -> Any:
                         ),
                     ],
                 )
-                if self.has_streaming and has_streaming != False
+                if self.has_streaming and has_streaming
                 else ChatOpenAI(
                     model_name=self.llm["model"],
                     openai_api_key=self._get_api_key(),
@@ -189,7 +189,7 @@ def _get_llm(self, has_streaming: bool = True) -> Any:
                         )
                     ],
                 )
-                if self.has_streaming and has_streaming != False
+                if self.has_streaming and has_streaming
                 else ChatAnthropic(anthropic_api_key=self._get_api_key())
             )
 
@@ -207,8 +207,10 @@ def _get_llm(self, has_streaming: bool = True) -> Any:
                         )
                     ],
                 )
-                if self.has_streaming and has_streaming != False
-                else Cohere(cohere_api_key=self._get_api_key(), model=self.llm["model"])
+                if self.has_streaming and has_streaming
+                else Cohere(
+                    cohere_api_key=self._get_api_key(), model=self.llm["model"]
+                )
             )
 
         if self.llm["provider"] == "azure-openai":
@@ -286,12 +288,10 @@ def _get_memory(self, session) -> ConversationBufferMemory:
         return ConversationBufferMemory(memory_key="chat_history", output_key="output")
 
     def _get_agent_documents(self) -> Any:
-        agent_documents = prisma.agentdocument.find_many(
+        return prisma.agentdocument.find_many(
             where={"agentId": self.id}, include={"document": True}
         )
 
-        return agent_documents
-
     def _get_tool_and_input_by_type(
         self, type: str, metadata: dict = None
     ) -> Tuple[Any, Any]:
@@ -407,12 +407,10 @@ def _get_tools(self) -> list[Tool]:
         return tools
 
     def _get_agent_tools(self) -> Any:
-        tools = prisma.agenttool.find_many(
+        return prisma.agenttool.find_many(
             where={"agentId": self.id}, include={"tool": True}
         )
 
-        return tools
-
     def _format_trace(self, trace: Any) -> dict:
         if self.documents or self.tools:
             return json.dumps(
@@ -463,7 +461,7 @@ def cache_message(
 
         metadata = {
             "agentId": agent_id,
-            "sessionId": str(session_id),
+            "sessionId": session_id,
             "text": query,
             "cached_message": ai_message,
             "type": "cache",
@@ -476,13 +474,11 @@ def cache_message(
 
     def get_cached_result(self, query: str) -> str | None:
         vectorstore: PineconeVectorStore = VectorStoreBase().get_database()
-        results = vectorstore.query(
+        if results := vectorstore.query(
             prompt=query,
             metadata_filter={"agentId": self.id},
             min_score=0.9,
-        )
-
-        if results:
+        ):
             timestamp: float = results[0].metadata.get("timestamp", 0.0)
 
             if timestamp and time.time() - timestamp > self.cache_ttl:

libs/legacy/app/lib/agents/factory.py

@@ -9,13 +9,8 @@ def create_agent(agent_base: AgentBase):
             if agent_base.tools or agent_base.documents:
                 return OpenAIAgent(agent_base)
 
-            return DefaultAgent(agent_base)
-
         elif agent_base.type == "REACT":
             if agent_base.tools or agent_base.documents:
                 return ReactAgent(agent_base)
 
-            return DefaultAgent(agent_base)
-
-        else:
-            return DefaultAgent(agent_base)
+        return DefaultAgent(agent_base)

libs/legacy/app/lib/api_tokens.py

@@ -2,5 +2,4 @@
 
 
 def generate_api_token() -> str:
-    api_token = uuid.uuid4().hex
-    return api_token
+    return uuid.uuid4().hex