Dev

[dcloud347]

This pull request introduces significant improvements to the handling of context and message data across multiple modules, centralizing logic for token management and enhancing type safety. The changes include creating a new CustomChatOpenAI class for better token handling, replacing str with a Context type in various states, and removing redundant truncation logic in favor of centralized token management.

Key Changes

Token Management and Model Enhancements:

• Introduced a new CustomChatOpenAI class in prometheus/chat_models/custom_chat_openai.py to handle token limits using tiktoken_counter, ensuring input messages are trimmed safely. This class is now used in LLMService to replace the previous ChatOpenAI integration. [1] [2] [3]

Context Type Migration:

• Replaced Sequence[str] with Sequence[Context] in BugReproductionState, ContextRetrievalState, and IssueClassificationState to improve type safety and semantic clarity. [1] [2] [3]
• Updated corresponding subgraphs and nodes to handle the new Context type, ensuring compatibility and proper formatting. [1] [2] [3]

Message Truncation Simplification:

• Removed the truncate_messages utility and its usage in several nodes (BugReproducingWriteNode, ContextProviderNode, EditNode, IssueBugAnalyzerNode) as token management is now centralized in CustomChatOpenAI. [1] [2] [3] [4]

String Conversion for Contexts:

• Ensured all contexts are explicitly converted to strings before joining in various nodes to handle the new Context type. [1] [2] [3] [4]

Code Cleanup:

• Removed unused imports and redundant retry logic in ContextRetrievalSubgraphNode, simplifying the codebase. [1] [2]

These changes collectively improve maintainability, enhance type safety, and centralize token management, making the codebase more robust and easier to extend.