Make flake8 happy

Author: Twixes

.github/workflows/ci.yml

@@ -36,7 +36,7 @@ jobs:
 
             - name: Lint with flake8
               run: |
-                  flake8 posthog --ignore E501
+                  flake8 posthog --ignore E501,W503
 
             - name: Check import order with isort
               run: |

posthog/ai/langchain/callbacks.py

@@ -19,21 +19,14 @@
 from uuid import UUID
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-    FunctionMessage,
-    HumanMessage,
-    SystemMessage,
-    ToolMessage,
-)
-from langchain_core.outputs import ChatGeneration, LLMResult
 from langchain.schema.agent import AgentAction, AgentFinish
+from langchain_core.messages import AIMessage, BaseMessage, FunctionMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
+from posthog import default_client
 from posthog.ai.utils import get_model_params, with_privacy_mode
 from posthog.client import Client
-from posthog import default_client
 
 log = logging.getLogger("posthog")
 

posthog/test/ai/langchain/test_callbacks.py

@@ -1,21 +1,20 @@
 import logging
 import math
 import os
-from pyexpat.errors import messages
 import time
-from typing import List, Optional, TypedDict, Union
 import uuid
-from unittest.mock import patch, ANY
+from typing import List, Optional, TypedDict, Union
+from unittest.mock import patch
 
 import pytest
 from langchain_anthropic.chat_models import ChatAnthropic
 from langchain_community.chat_models.fake import FakeMessagesListChatModel
 from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnableLambda
 from langchain_openai.chat_models import ChatOpenAI
-from langgraph.graph.state import StateGraph, START, END
+from langgraph.graph.state import END, START, StateGraph
 
 from posthog.ai.langchain import CallbackHandler