add support for azure backend (#88)

## Description

<!--
Please include a summary of the changes below;
Fill in the issue number that this PR addresses (if applicable);
Mention the person who will review this PR (if you know who it is);
Replace (summary), (issue), and (reviewer) with the appropriate
information (No parentheses).

请在下方填写更改的摘要；
填写此 PR 解决的问题编号（如果适用）；
提及将审查此 PR 的人（如果您知道是谁）；
替换 (summary)、(issue) 和 (reviewer) 为适当的信息（不带括号）。
-->

Summary: (summary) Add azure backend support

Fix: #(issue)

Reviewer: @(reviewer)

## Checklist:

- [ x] I have performed a self-review of my own code | 我已自行检查了自己的代码
- [x ] I have commented my code in hard-to-understand areas |
我已在难以理解的地方对代码进行了注释
- [x ] I have added tests that prove my fix is effective or that my
feature works | 我已添加测试以证明我的修复有效或功能正常
- [ ] I have added necessary documentation (if applicable) |
我已添加必要的文档（如果适用）
- [ ] I have linked the issue to this PR (if applicable) | 我已将 issue
链接到此 PR（如果适用）
- [ ] I have mentioned the person who will review this PR | 我已提及将审查此 PR
的人

Author: tangg555

src/memos/configs/llm.py

@@ -26,6 +26,18 @@ class OpenAILLMConfig(BaseLLMConfig):
     )
 
 
+class AzureLLMConfig(BaseLLMConfig):
+    base_url: str = Field(
+        default="https://api.openai.azure.com/",
+        description="Base URL for Azure OpenAI API",
+    )
+    api_version: str = Field(
+        default="2024-03-01-preview",
+        description="API version for Azure OpenAI",
+    )
+    api_key: str = Field(..., description="API key for Azure OpenAI")
+
+
 class OllamaLLMConfig(BaseLLMConfig):
     api_base: str = Field(
         default="http://localhost:11434",
@@ -61,6 +73,7 @@ class LLMConfigFactory(BaseConfig):
     backend_to_class: ClassVar[dict[str, Any]] = {
         "openai": OpenAILLMConfig,
         "ollama": OllamaLLMConfig,
+        "azure": AzureLLMConfig,
         "huggingface": HFLLMConfig,
         "vllm": VLLMLLMConfig,
         "huggingface_singleton": HFLLMConfig,  # Add singleton support

src/memos/llms/factory.py

@@ -5,7 +5,7 @@
 from memos.llms.hf import HFLLM
 from memos.llms.hf_singleton import HFSingletonLLM
 from memos.llms.ollama import OllamaLLM
-from memos.llms.openai import OpenAILLM
+from memos.llms.openai import AzureLLM, OpenAILLM
 from memos.llms.vllm import VLLMLLM
 
 
@@ -14,6 +14,7 @@ class LLMFactory(BaseLLM):
 
     backend_to_class: ClassVar[dict[str, Any]] = {
         "openai": OpenAILLM,
+        "azure": AzureLLM,
         "ollama": OllamaLLM,
         "huggingface": HFLLM,
         "huggingface_singleton": HFSingletonLLM,  # Add singleton version

src/memos/llms/openai.py

@@ -1,6 +1,6 @@
 import openai
 
-from memos.configs.llm import OpenAILLMConfig
+from memos.configs.llm import AzureLLMConfig, OpenAILLMConfig
 from memos.llms.base import BaseLLM
 from memos.llms.utils import remove_thinking_tags
 from memos.log import get_logger
@@ -32,3 +32,31 @@ def generate(self, messages: MessageList) -> str:
             return remove_thinking_tags(response_content)
         else:
             return response_content
+
+
+class AzureLLM(BaseLLM):
+    """Azure OpenAI LLM class."""
+
+    def __init__(self, config: AzureLLMConfig):
+        self.config = config
+        self.client = openai.AzureOpenAI(
+            azure_endpoint=config.base_url,
+            api_version=config.api_version,
+            api_key=config.api_key,
+        )
+
+    def generate(self, messages: MessageList) -> str:
+        """Generate a response from Azure OpenAI LLM."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+        )
+        logger.info(f"Response from Azure OpenAI: {response.model_dump_json()}")
+        response_content = response.choices[0].message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        else:
+            return response_content

src/memos/memories/textual/general.py

@@ -8,7 +8,7 @@
 
 from memos.configs.memory import GeneralTextMemoryConfig
 from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
-from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM
+from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM, AzureLLM
 from memos.log import get_logger
 from memos.memories.textual.base import BaseTextMemory
 from memos.memories.textual.item import TextualMemoryItem
@@ -26,7 +26,9 @@ class GeneralTextMemory(BaseTextMemory):
     def __init__(self, config: GeneralTextMemoryConfig):
         """Initialize memory with the given configuration."""
         self.config: GeneralTextMemoryConfig = config
-        self.extractor_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.extractor_llm)
+        self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.extractor_llm
+        )
         self.vector_db: QdrantVecDB = VecDBFactory.from_config(config.vector_db)
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
 

src/memos/memories/textual/tree.py

@@ -10,7 +10,7 @@
 from memos.configs.memory import TreeTextMemoryConfig
 from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
 from memos.graph_dbs.factory import GraphStoreFactory, Neo4jGraphDB
-from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM
+from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM, AzureLLM
 from memos.log import get_logger
 from memos.memories.textual.base import BaseTextMemory
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
@@ -31,8 +31,12 @@ class TreeTextMemory(BaseTextMemory):
     def __init__(self, config: TreeTextMemoryConfig):
         """Initialize memory with the given configuration."""
         self.config: TreeTextMemoryConfig = config
-        self.extractor_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.extractor_llm)
-        self.dispatcher_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.dispatcher_llm)
+        self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.extractor_llm
+        )
+        self.dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.dispatcher_llm
+        )
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
         self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
         self.is_reorganize = config.reorganize

src/memos/memories/textual/tree_text_memory/organize/manager.py

@@ -5,7 +5,7 @@
 
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.neo4j import Neo4jGraphDB
-from memos.llms.factory import OllamaLLM, OpenAILLM
+from memos.llms.factory import OllamaLLM, OpenAILLM, AzureLLM
 from memos.log import get_logger
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
 from memos.memories.textual.tree_text_memory.organize.reorganizer import (
@@ -22,7 +22,7 @@ def __init__(
         self,
         graph_store: Neo4jGraphDB,
         embedder: OllamaEmbedder,
-        llm: OpenAILLM | OllamaLLM,
+        llm: OpenAILLM | OllamaLLM | AzureLLM,
         memory_size: dict | None = None,
         threshold: float | None = 0.80,
         merged_threshold: float | None = 0.92,

src/memos/memories/textual/tree_text_memory/retrieve/reranker.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from memos.embedders.factory import OllamaEmbedder
-from memos.llms.factory import OllamaLLM, OpenAILLM
+from memos.llms.factory import OllamaLLM, OpenAILLM, AzureLLM
 from memos.memories.textual.item import TextualMemoryItem
 from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
 
@@ -41,7 +41,7 @@ class MemoryReranker:
     Rank retrieved memory cards by structural priority and contextual similarity.
     """
 
-    def __init__(self, llm: OpenAILLM | OllamaLLM, embedder: OllamaEmbedder):
+    def __init__(self, llm: OpenAILLM | OllamaLLM | AzureLLM, embedder: OllamaEmbedder):
         self.llm = llm
         self.embedder = embedder
 

src/memos/memories/textual/tree_text_memory/retrieve/searcher.py

@@ -5,7 +5,7 @@
 
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.factory import Neo4jGraphDB
-from memos.llms.factory import OllamaLLM, OpenAILLM
+from memos.llms.factory import OllamaLLM, OpenAILLM, AzureLLM
 from memos.memories.textual.item import SearchedTreeNodeTextualMemoryMetadata, TextualMemoryItem
 
 from .internet_retriever_factory import InternetRetrieverFactory
@@ -18,7 +18,7 @@
 class Searcher:
     def __init__(
         self,
-        dispatcher_llm: OpenAILLM | OllamaLLM,
+        dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM,
         graph_store: Neo4jGraphDB,
         embedder: OllamaEmbedder,
         internet_retriever: InternetRetrieverFactory | None = None,