Buckets:
tahamajs/Sysmem2_in_AI / ComputerAssignments /CA15_test_time_training /src /langchain_integration.py
| import os | |
| import json | |
| from typing import List, Dict, Any, Optional, Union | |
| from dataclasses import dataclass | |
| from enum import Enum | |
| try: | |
| from langchain.llms import OpenAI, GooglePalm | |
| from langchain.chat_models import ChatOpenAI, ChatGoogleGenerativeAI | |
| from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage | |
| from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory | |
| from langchain.chains import LLMChain, SequentialChain, TransformChain | |
| from langchain.prompts import PromptTemplate, ChatPromptTemplate | |
| from langchain.tools import BaseTool, Tool | |
| from langchain.agents import AgentExecutor, create_react_agent, initialize_agent | |
| from langchain.agents.agent_types import AgentType | |
| from langchain.callbacks import BaseCallbackHandler | |
| from langchain.schema import AgentAction, AgentFinish | |
| LANCHAIN_AVAILABLE = True | |
| except ImportError: | |
| LANCHAIN_AVAILABLE = False | |
| print("⚠️ LangChain not available. Install with: pip install langchain") | |
| try: | |
| import google.generativeai as genai | |
| GEMINI_AVAILABLE = True | |
| except ImportError: | |
| GEMINI_AVAILABLE = False | |
| print("⚠️ Google Generative AI not available. Install with: pip install google-generativeai") | |
| class TTTTaskType(Enum): | |
| PATTERN_RECOGNITION = "pattern_recognition" | |
| TEXT_TRANSFORMATION = "text_transformation" | |
| REASONING_TASK = "reasoning_task" | |
| CREATIVE_WRITING = "creative_writing" | |
| CODE_GENERATION = "code_generation" | |
| MULTI_STEP_REASONING = "multi_step_reasoning" | |
| class TTTTaskConfig: | |
| task_type: TTTTaskType | |
| description: str | |
| examples: List[Dict[str, str]] | |
| complexity_level: int | |
| required_tools: List[str] | |
| memory_type: str = "buffer" | |
| class TTTMemoryManager: | |
| def __init__(self, memory_type: str = "buffer", max_token_limit: int = 2000): | |
| self.memory_type = memory_type | |
| self.max_token_limit = max_token_limit | |
| self.memories = {} | |
| def get_memory(self, session_id: str) -> Union[ConversationBufferMemory, ConversationSummaryMemory]: | |
| if session_id not in self.memories: | |
| if self.memory_type == "buffer": | |
| self.memories[session_id] = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| return_messages=True, | |
| max_token_limit=self.max_token_limit | |
| ) | |
| elif self.memory_type == "summary": | |
| self.memories[session_id] = ConversationSummaryMemory( | |
| llm=ChatOpenAI(temperature=0), | |
| memory_key="chat_history", | |
| return_messages=True, | |
| max_token_limit=self.max_token_limit | |
| ) | |
| return self.memories[session_id] | |
| def save_context(self, session_id: str, input_text: str, output_text: str): | |
| memory = self.get_memory(session_id) | |
| memory.save_context({"input": input_text}, {"output": output_text}) | |
| def get_context(self, session_id: str) -> str: | |
| memory = self.get_memory(session_id) | |
| return memory.buffer | |
| class TTTToolKit: | |
| def __init__(self): | |
| self.tools = self._create_tools() | |
| def _create_tools(self) -> List[Tool]: | |
| tools = [] | |
| def analyze_pattern(input_text: str) -> str: | |
| patterns = [] | |
| if len(input_text) > 1: | |
| patterns.append(f"Length: {len(input_text)}") | |
| patterns.append(f"First char: {input_text[0]}") | |
| patterns.append(f"Last char: {input_text[-1]}") | |
| patterns.append(f"Reversed: {input_text[::-1]}") | |
| return " | ".join(patterns) | |
| tools.append(Tool( | |
| name="pattern_analyzer", | |
| description="Analyze text patterns for TTT tasks", | |
| func=analyze_pattern | |
| )) | |
| def transform_text(input_text: str, transformation: str) -> str: | |
| if transformation == "reverse": | |
| return input_text[::-1] | |
| elif transformation == "uppercase": | |
| return input_text.upper() | |
| elif transformation == "lowercase": | |
| return input_text.lower() | |
| elif transformation == "reverse_append_last": | |
| return input_text[::-1] + input_text[-1] | |
| else: | |
| return input_text | |
| tools.append(Tool( | |
| name="text_transformer", | |
| description="Transform text using various methods", | |
| func=transform_text | |
| )) | |
| def calculate_complexity(input_text: str) -> str: | |
| complexity_score = 0 | |
| complexity_score += len(input_text) * 0.1 | |
| complexity_score += len(set(input_text)) * 0.2 | |
| if any(c.isdigit() for c in input_text): | |
| complexity_score += 1 | |
| if any(c in "!@#$%^&*()" for c in input_text): | |
| complexity_score += 1 | |
| if complexity_score < 2: | |
| return "Low complexity" | |
| elif complexity_score < 4: | |
| return "Medium complexity" | |
| else: | |
| return "High complexity" | |
| tools.append(Tool( | |
| name="complexity_calculator", | |
| description="Calculate task complexity", | |
| func=calculate_complexity | |
| )) | |
| return tools | |
| class TTTChainBuilder: | |
| def __init__(self, llm_model: str = "gpt-3.5-turbo", api_key: Optional[str] = None): | |
| self.llm_model = llm_model | |
| self.api_key = api_key or os.getenv("OPENAI_API_KEY") | |
| self.llm = self._initialize_llm() | |
| self.toolkit = TTTToolKit() | |
| def _initialize_llm(self): | |
| if not LANCHAIN_AVAILABLE: | |
| return None | |
| if self.llm_model.startswith("gpt"): | |
| return ChatOpenAI( | |
| model_name=self.llm_model, | |
| temperature=0.7, | |
| openai_api_key=self.api_key | |
| ) | |
| elif self.llm_model.startswith("gemini"): | |
| return ChatGoogleGenerativeAI( | |
| model=self.llm_model, | |
| temperature=0.7, | |
| google_api_key=os.getenv("GOOGLE_API_KEY") | |
| ) | |
| else: | |
| return ChatOpenAI(temperature=0.7) | |
| def create_pattern_recognition_chain(self) -> LLMChain: | |
| if not self.llm: | |
| return None | |
| prompt_template = PromptTemplate( | |
| input_variables=["examples", "query", "analysis"], | |
| template= | |
| ) | |
| return LLMChain(llm=self.llm, prompt=prompt_template) | |
| def create_reasoning_chain(self) -> LLMChain: | |
| if not self.llm: | |
| return None | |
| prompt_template = PromptTemplate( | |
| input_variables=["task_description", "examples", "query"], | |
| template= | |
| ) | |
| return LLMChain(llm=self.llm, prompt=prompt_template) | |
| def create_multi_step_chain(self) -> SequentialChain: | |
| if not self.llm: | |
| return None | |
| analysis_prompt = PromptTemplate( | |
| input_variables=["input_text"], | |
| template="Analyze this text: {input_text}. What patterns do you see?" | |
| ) | |
| analysis_chain = LLMChain(llm=self.llm, prompt=analysis_prompt, output_key="analysis") | |
| transform_prompt = PromptTemplate( | |
| input_variables=["input_text", "analysis"], | |
| template="Based on this analysis: {analysis}, transform this text: {input_text}" | |
| ) | |
| transform_chain = LLMChain(llm=self.llm, prompt=transform_prompt, output_key="transformation") | |
| verify_prompt = PromptTemplate( | |
| input_variables=["input_text", "transformation"], | |
| template="Verify this transformation: {input_text} -> {transformation}. Is it correct?" | |
| ) | |
| verify_chain = LLMChain(llm=self.llm, prompt=verify_prompt, output_key="verification") | |
| return SequentialChain( | |
| chains=[analysis_chain, transform_chain, verify_chain], | |
| input_variables=["input_text"], | |
| output_variables=["analysis", "transformation", "verification"] | |
| ) | |
| class TTTAgent: | |
| def __init__(self, llm_model: str = "gpt-3.5-turbo", api_key: Optional[str] = None): | |
| self.chain_builder = TTTChainBuilder(llm_model, api_key) | |
| self.toolkit = TTTToolKit() | |
| self.memory_manager = TTTMemoryManager() | |
| self.agent = None | |
| self._initialize_agent() | |
| def _initialize_agent(self): | |
| if not LANCHAIN_AVAILABLE or not self.chain_builder.llm: | |
| return | |
| agent_prompt = PromptTemplate( | |
| input_variables=["input", "agent_scratchpad", "tools", "tool_names"], | |
| template= | |
| ) | |
| self.agent = create_react_agent( | |
| llm=self.chain_builder.llm, | |
| tools=self.toolkit.tools, | |
| prompt=agent_prompt | |
| ) | |
| def process_ttt_task(self, task_config: TTTTaskConfig, query: str, session_id: str = "default") -> Dict[str, Any]: | |
| if not self.agent: | |
| return {"error": "Agent not initialized"} | |
| memory = self.memory_manager.get_memory(session_id) | |
| examples_str = "\n".join([f"{ex['input']} -> {ex['output']}" for ex in task_config.examples]) | |
| agent_input = { | |
| "input": f | |
| "agent_scratchpad": "", | |
| "tools": [tool.description for tool in self.toolkit.tools], | |
| "tool_names": [tool.name for tool in self.toolkit.tools] | |
| } | |
| try: | |
| agent_executor = AgentExecutor( | |
| agent=self.agent, | |
| tools=self.toolkit.tools, | |
| memory=memory, | |
| verbose=True, | |
| handle_parsing_errors=True | |
| ) | |
| result = agent_executor.run(agent_input) | |
| self.memory_manager.save_context(session_id, query, result) | |
| return { | |
| "success": True, | |
| "result": result, | |
| "task_type": task_config.task_type.value, | |
| "session_id": session_id | |
| } | |
| except Exception as e: | |
| return { | |
| "success": False, | |
| "error": str(e), | |
| "task_type": task_config.task_type.value | |
| } | |
| class GeminiTTTIntegration: | |
| def __init__(self, api_key: Optional[str] = None): | |
| self.api_key = api_key or os.getenv("GOOGLE_API_KEY") | |
| self.model = None | |
| self._initialize_gemini() | |
| def _initialize_gemini(self): | |
| if not GEMINI_AVAILABLE or not self.api_key: | |
| print("⚠️ Gemini API not available") | |
| return | |
| try: | |
| genai.configure(api_key=self.api_key) | |
| self.model = genai.GenerativeModel('gemini-pro') | |
| print("✅ Gemini API initialized successfully") | |
| except Exception as e: | |
| print(f"❌ Failed to initialize Gemini: {e}") | |
| def generate_creative_examples(self, task_description: str, num_examples: int = 5) -> List[Dict[str, str]]: | |
| if not self.model: | |
| return [] | |
| prompt = f | |
| try: | |
| response = self.model.generate_content(prompt) | |
| examples = json.loads(response.text) | |
| return examples[:num_examples] | |
| except Exception as e: | |
| print(f"❌ Failed to generate examples: {e}") | |
| return [] | |
| def evaluate_ttt_performance(self, examples: List[Dict[str, str]], predictions: List[str]) -> Dict[str, Any]: | |
| if not self.model: | |
| return {"error": "Gemini not available"} | |
| prompt = f | |
| try: | |
| response = self.model.generate_content(prompt) | |
| evaluation = json.loads(response.text) | |
| return evaluation | |
| except Exception as e: | |
| return {"error": f"Evaluation failed: {e}"} | |
| def generate_task_variations(self, base_task: str, num_variations: int = 3) -> List[str]: | |
| if not self.model: | |
| return [] | |
| prompt = f | |
| try: | |
| response = self.model.generate_content(prompt) | |
| variations = json.loads(response.text) | |
| return variations[:num_variations] | |
| except Exception as e: | |
| print(f"❌ Failed to generate variations: {e}") | |
| return [] | |
| class AdvancedTTTWorkflow: | |
| def __init__(self, llm_model: str = "gpt-3.5-turbo", gemini_api_key: Optional[str] = None): | |
| self.ttt_agent = TTTAgent(llm_model) | |
| self.gemini_integration = GeminiTTTIntegration(gemini_api_key) | |
| self.chain_builder = TTTChainBuilder(llm_model) | |
| def run_advanced_ttt_experiment(self, task_config: TTTTaskConfig, test_queries: List[str]) -> Dict[str, Any]: | |
| results = { | |
| "task_config": task_config, | |
| "test_queries": test_queries, | |
| "agent_results": [], | |
| "gemini_evaluation": None, | |
| "creative_examples": [], | |
| "task_variations": [] | |
| } | |
| if self.gemini_integration.model: | |
| creative_examples = self.gemini_integration.generate_creative_examples( | |
| task_config.description, 5 | |
| ) | |
| results["creative_examples"] = creative_examples | |
| if self.gemini_integration.model: | |
| variations = self.gemini_integration.generate_task_variations( | |
| task_config.description, 3 | |
| ) | |
| results["task_variations"] = variations | |
| predictions = [] | |
| for query in test_queries: | |
| agent_result = self.ttt_agent.process_ttt_task(task_config, query) | |
| results["agent_results"].append(agent_result) | |
| if agent_result.get("success"): | |
| predictions.append(agent_result["result"]) | |
| else: | |
| predictions.append("ERROR") | |
| if self.gemini_integration.model and predictions: | |
| evaluation = self.gemini_integration.evaluate_ttt_performance( | |
| task_config.examples, predictions | |
| ) | |
| results["gemini_evaluation"] = evaluation | |
| return results | |
| def create_sample_task_configs() -> List[TTTTaskConfig]: | |
| configs = [] | |
| configs.append(TTTTaskConfig( | |
| task_type=TTTTaskType.PATTERN_RECOGNITION, | |
| description="Reverse the word and append the last character", | |
| examples=[ | |
| {"input": "hello", "output": "olleho"}, | |
| {"input": "world", "output": "dlrowd"}, | |
| {"input": "python", "output": "nohtypp"} | |
| ], | |
| complexity_level=2, | |
| required_tools=["pattern_analyzer", "text_transformer"] | |
| )) | |
| configs.append(TTTTaskConfig( | |
| task_type=TTTTaskType.REASONING_TASK, | |
| description="Apply Fibonacci sequence transformation to letters", | |
| examples=[ | |
| {"input": "abc", "output": "bdf"}, | |
| {"input": "xyz", "output": "yza"}, | |
| {"input": "test", "output": "ufuv"} | |
| ], | |
| complexity_level=4, | |
| required_tools=["pattern_analyzer", "complexity_calculator"] | |
| )) | |
| return configs | |
| if __name__ == "__main__": | |
| print("🧠 Testing LangChain Integration for TTT") | |
| print("=" * 50) | |
| task_config = create_sample_task_configs()[0] | |
| print(f"📋 Task: {task_config.description}") | |
| workflow = AdvancedTTTWorkflow() | |
| test_queries = ["gemini", "tensor", "neural"] | |
| print("🚀 Running advanced TTT experiment...") | |
| results = workflow.run_advanced_ttt_experiment(task_config, test_queries) | |
| print("📊 Results:") | |
| print(f" Agent Results: {len(results['agent_results'])}") | |
| print(f" Creative Examples: {len(results['creative_examples'])}") | |
| print(f" Task Variations: {len(results['task_variations'])}") | |
| print(f" Gemini Evaluation: {results['gemini_evaluation'] is not None}") | |
| print("✅ LangChain integration test completed!") | |
Xet Storage Details
- Size:
- 16.2 kB
- Xet hash:
- fdf920311a968f6178309fcb17c3bd14b7e51f85699039eeee8905e5d3de64e9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.