Trying to integrate LightRAG into AutoGen

[PatrikPatrik2]

I am trying to integrate LigthRag in AutoGen, but I am not sucessfull. I can get it answer the RAG question, but either it gets into an infinit loop, or if I add max_consecutive_auto_reply=1 it only do it once, but then I get this error:

`Studsvik_rag_agent (to assistant):

Conversation complete.

---

Error in main: 'AzureOpenAI' object has no attribute 'total_usage_summary' `

I have got completely stuck on this - Does anyone see what I have done wrong?

import asyncio  
import nest_asyncio  
import numpy as np  
from autogen import ConversableAgent  
from lightrag import LightRAG, QueryParam  
from lightrag.utils import EmbeddingFunc  
from openai import AzureOpenAI  
from functools import partial  
from config import settings  # Import settings from config/settings.py  

# Apply nest_asyncio to allow nested event loops  
nest_asyncio.apply()  

class RAGAgent(ConversableAgent):  
    def __init__(self, name, working_dir="./RagStudsvik", **kwargs):  
        super().__init__(name, **kwargs)  

        # Initialize Azure OpenAI client  
        self.client = AzureOpenAI(  
            api_key=settings.AZURE_OPENAI_API_KEY,  
            api_version="2023-05-15",  
            azure_endpoint=settings.AZURE_OPENAI_ENDPOINT  
        )  

        # Store working directory  
        self.working_dir = working_dir  
        self.message_history = []  
        self.conversation_started = False  

        # Initialize RAG with bound methods  
        self.rag = self._create_rag()  

    def llm_model_func(self, prompt, system_prompt=None, history_messages=[], **kwargs):  
        messages = []  
        if system_prompt:  
            messages.append({"role": "system", "content": system_prompt})  
        messages.extend(history_messages)  
        messages.append({"role": "user", "content": prompt})  

        kwargs.pop('hashing_kv', None)  
        try:  
            response = self.client.chat.completions.create(  
                model="gpt4o",  
                messages=messages,  
                temperature=0,  
                **kwargs  
            )  
            return response.choices[0].message.content  
        except Exception as e:  
            print(f"Error in llm_model_func: {str(e)}")  
            return None  

    def embedding_func(self, texts: list[str]) -> np.ndarray:  
        try:  
            embeddings = []  
            for text in texts:  
                response = self.client.embeddings.create(  
                    model="embedding",  
                    input=text  
                )  
                embeddings.append(response.data[0].embedding)  
            return np.array(embeddings)  
        except Exception as e:  
            print(f"Error in embedding_func: {str(e)}")  
            return None  

    def _create_rag(self):  
        try:  
            async def async_llm_wrapper(*args, **kwargs):  
                return self.llm_model_func(*args, **kwargs)  

            async def async_embedding_wrapper(*args, **kwargs):  
                return self.embedding_func(*args, **kwargs)  

            return LightRAG(  
                working_dir=self.working_dir,  
                llm_model_func=async_llm_wrapper,  
                embedding_func=EmbeddingFunc(  
                    embedding_dim=3072,  
                    max_token_size=8192,  
                    func=async_embedding_wrapper  
                ),  
                tiktoken_model_name="gpt-4",  
                llm_model_name="gpt-4"  
            )  
        except Exception as e:  
            print(f"Error initializing RAG: {str(e)}")  
            raise  

    async def _process_message(self, message):  
        try:  
            self.message_history.append({"role": "user", "content": message})  
            result = await self.rag.aquery(  
                message,   
                param=QueryParam(mode="hybrid")  
            )  
            self.message_history.append({"role": "assistant", "content": result})  
            return result  
        except Exception as e:  
            print(f"Error processing message: {str(e)}")  
            return f"Error processing message: {str(e)}"  

    async def a_generate_reply(self, messages=None, sender=None, **kwargs):  
        """Async method for generating replies with loop control"""  
        if not messages:  
            return None  

        # Check if this is a new conversation  
        if not self.conversation_started:  
            self.conversation_started = True  
            last_message = messages[-1]  
            if isinstance(last_message, dict):  
                message_content = last_message.get("content", "")  
            else:  
                message_content = str(last_message)  

            try:  
                reply = await self._process_message(message_content)  
                # Add a termination signal  
                return {"content": reply, "terminate": True}  
            except Exception as e:  
                print(f"Error generating reply: {str(e)}")  
                return {"content": f"Error generating reply: {str(e)}", "terminate": True}  
        else:  
            # For subsequent messages, indicate that the conversation should end  
            return {"content": "Conversation complete.", "terminate": True}  

    def generate_reply(self, messages=None, sender=None, **kwargs):  
        try:  
            return asyncio.run(self.a_generate_reply(messages, sender, **kwargs))  
        except Exception as e:  
            print(f"Error in generate_reply: {str(e)}")  
            return {"content": f"Error in generate_reply: {str(e)}", "terminate": True}  

def create_rag_agent(name="Studsvik_rag_agent", working_dir="./RagStudsvik", system_message=None):  
    try:  
        # Read the system message from the file  

        # Create the RAG agent  
        agent = RAGAgent(  
            name=name,  
            working_dir=working_dir,  
            #system_message=system_message,  
            system_message="You are a helpful RAG agent",  
            llm_config={  
                "config_list": [{  
                    "model": "gpt4o",  
                    "api_key": settings.AZURE_OPENAI_API_KEY,  
                    "base_url": settings.AZURE_OPENAI_ENDPOINT,  
                    "api_type": "azure",  
                    "api_version": "2023-05-15"  
                }]  
            },  
            human_input_mode="NEVER",
            max_consecutive_auto_reply=1   
        )  
        return agent  
    except Exception as e:  
        print(f"Error creating RAG agent: {str(e)}")  
        raise  

if __name__ == "__main__":  
    try:  
        # Create the RAG agent  
        rag_agent = create_rag_agent()  

        with open("Studsvik_rag_agent_system_message.txt", "r") as file:  
            system_message_Studsvik = file.read()

        # Create another agent to interact with the RAG agent  
        assistant = ConversableAgent(  
            "assistant",  
            system_message=system_message_Studsvik,  
            llm_config={  
                "config_list": [{  
                    "model": "gpt4o",  
                    "api_key": settings.AZURE_OPENAI_API_KEY,  
                    "base_url": settings.AZURE_OPENAI_ENDPOINT,  
                    "api_type": "azure",  
                    "api_version": "2023-05-15"  
                }]  
            },  
            human_input_mode="NEVER",
            max_consecutive_auto_reply=1    
        )  

        # Start a conversation with a single question  
        assistant.initiate_chat(  
            rag_agent,  
            message="What is optimal  formulation?"  
        )  
    except Exception as e:  
        print(f"Error in main: {str(e)}")  

[rysweet]

hi @PatrikPatrik2 - thanks for the bug report. Its not super clear to me from the post how one would replicate this, and it looks like the error you posted "Error in main: 'AzureOpenAI' object has no attribute 'total_usage_summary'" is incomplete - could you post the whole stack trace?

[ekzhu]

@PatrikPatrik2 please share your installed autogen package, version, and used model name. Also the full output.

Have you looked into the v0.4 preview API? I think your issue is resolved in the latest version.