poe-langchain. Adapting LLM that we can access on `poe`, into langchain's `llm` eco-system.

poe-langchain.py

from langchain_core.language_models.llms import BaseLLM
from poe_api_wrapper import PoeApi
from pydantic import BaseModel, Field

class PoeApiAdapter(BaseLLM):
    tokens: dict
    bot: str
    auto_proxy: bool = Field(default=False)
    
    class Config:
        extra = "allow"  # Allows assignment of attributes not declared in the model

    def __init__(self, **data):
        super().__init__(**data)
        self._initialize_client()

    def _initialize_client(self):
        self.client = PoeApi(cookie=self.tokens, auto_proxy=False)

    def __call__(self, prompt, stop=None, callbacks=None, *, tags=None, metadata=None, **kwargs):
        # This method sends a message and waits for a complete response
        return self._generate(prompt, stop, callbacks, tags=tags, metadata=metadata, **kwargs)

    def invoke(self, prompt, **kwargs):
        # This method sends a message and waits for a complete response
        final_response = None
        for chunk in self.client.send_message(self.bot, prompt):
            if chunk['state'] == 'complete':
                final_response = chunk
                break
        return final_response['text'] if final_response else "No complete response received."

    def _generate(self, prompt, stop=None, callbacks=None, tags=None, metadata=None, **kwargs):
        final_response = None
        for chunk in self.client.send_message(self.bot, prompt):
            if chunk['state'] == 'complete':
                final_response = chunk
                break
        response_text = final_response['text'] if final_response else "No complete response received."
        if self.verbose:
            print(response_text)
        return response_text

    def _llm_type(self):
        # Provide a unique identifier for this LLM type
        return "PoeApiBot"

      
# Example usage
import os

# Setup tokens and bot
tokens = {
    'b': os.getenv('poe-p-b-cookie'),
    'lat': os.getenv('poe-p-lat-cookie')
}
bot = "a2"  # Specify the bot you want to use

# Create an instance of your adapter
llm = PoeApiAdapter(tokens=tokens, bot=bot)

# Now, use it in a Langchain context using invoke, which internally will use the generate method
response = llm.invoke("What is reverse engineering?")
print(response)