Merge pull request #9 from egehanyorulmaz/dev/improve_company_metadata

Refactor: Parallelize Company Search with a Dynamic Graph

Author: egehanyorulmaz

src/core/agents/graph_builder.py

@@ -26,14 +26,29 @@ def build_resume_analysis_graph():
 
     # Use functools.partial to inject dependencies into the nodes
     workflow.add_node("parse_resume", partial(parse_resume_node, extractor=llm_extractor))
-    workflow.add_node("company_research", partial(search_company_info_node, template_service=template_service, model_name="gpt-4o"))
+    workflow.add_node("job_company_research", partial(search_company_info_node, 
+                                                      template_service=template_service, 
+                                                      branch="job_description",
+                                                      model_name="gpt-4.1-mini"))
+    workflow.add_node("resume_company_research", partial(search_company_info_node, 
+                                                         template_service=template_service, 
+                                                         branch="resume",
+                                                         model_name="gpt-4.1-mini"))
+
     workflow.add_node("parse_job_description", partial(parse_job_description_node, extractor=llm_extractor))
     workflow.add_node("experience_analyzer", partial(analyze_experience_node, extractor=llm_extractor))
-    # Connect nodes sequentially
+    
+    #### RESUME PATH ###
     workflow.add_edge(START, "parse_resume")
-    workflow.add_edge("parse_resume", "company_research")
-    workflow.add_edge("company_research", "parse_job_description")
-    workflow.add_edge("parse_job_description", "experience_analyzer")
+    workflow.add_edge("parse_resume", "resume_company_research")
+    
+    #### JOB DESCRIPTION PATH ###
+    workflow.add_edge(START, "parse_job_description")
+    workflow.add_edge("parse_job_description", "job_company_research")
+
+    #### ANALYSIS PATH ###
+    workflow.add_edge("resume_company_research", "experience_analyzer")
+    workflow.add_edge("job_company_research", "experience_analyzer")
     workflow.add_edge("experience_analyzer", END)
 
     # Compile and return the graph

src/core/agents/search_agents.py

@@ -4,12 +4,14 @@
 from langchain_openai import ChatOpenAI
 
 from pydantic import BaseModel, Field
-from typing import Type, Optional, Dict, Any, List
+from typing import Type, Optional, Dict, Any, List, Literal
+import logging
 
 from src.core.ports.secondary.template_service import TemplateService
 from src.core.agents.utils.state import AgentState
-from src.core.domain.company_search import CompanySearchResponse
+from src.core.domain.company_search import CompanyInfo
 
+logger = logging.getLogger("core.agents.search_agents")
 
 def get_search_tool(max_results=3, time_range="month"):
     """
@@ -30,7 +32,7 @@ def get_search_tool(max_results=3, time_range="month"):
         topic="general",           # Use "general" for company information
     )
 
-async def create_company_search_agent(response_format: Type[BaseModel], 
+def create_company_search_agent(response_format: Type[BaseModel], 
                                 template_service: TemplateService,
                                 model_name="gpt-3.5-turbo"):
     """
@@ -65,25 +67,26 @@ async def create_company_search_agent(response_format: Type[BaseModel],
 
 async def search_company_info(company_name: str, 
                         template_service: TemplateService, 
-                        model_name="gpt-3.5-turbo") -> Optional[CompanySearchResponse]:
+                        model_name="gpt-3.5-turbo") -> Optional[CompanyInfo]:
     """
     Search for up-to-date information about a company and return structured data.
     """
+    logger.info(f"Collecting more details on {company_name}")
     # TODO: Retry on OpenAI rate limit errors
     search_query = template_service.render_prompt(
         "prompts/company_search/search_query.j2",
-        **{"company_name": company_name, "search_result_format": CompanySearchResponse.model_json_schema()}
+        **{"company_name": company_name, "search_result_format": CompanyInfo.model_json_schema()}
     )
 
     try:
-        agent = await create_company_search_agent(CompanySearchResponse, template_service, model_name)
+        agent = create_company_search_agent(CompanyInfo, template_service, model_name)
         response = await agent.ainvoke(
             {"messages": [{"role": "user", "content": search_query}]}
         )
 
         if "structured_response" in response:
             result = response["structured_response"]
-            return CompanySearchResponse.model_validate(result)
+            return CompanyInfo.model_validate(result)
         else:
             print("No structured response found in agent output")
             return None
@@ -92,23 +95,39 @@ async def search_company_info(company_name: str,
         return None
 
 
-async def main():
+
+
+if __name__ == "__main__":
+    import asyncio
     from src.core.domain.config import TemplateConfig
     from src.infrastructure.template.jinja_template_service import JinjaTemplateService
     template_config = TemplateConfig.development()
     template_service = JinjaTemplateService(config=template_config)
 
-    company_info = await search_company_info("Apple Inc.", template_service, model_name="gpt-4o")
-    
-    if company_info:
-        print("\nCompany Information:")
-        print(f"Name: {company_info.company_name}")
-        print(f"Industry: {company_info.company_industry}")
-        print(f"Size: {company_info.company_size} employees")
-        print(f"Revenue: {company_info.company_revenue}")
-        print(f"Location: {company_info.company_location}")
-        print(f"Website: {company_info.company_website}")
-        print(f"Founded: {company_info.founded_year or 'Unknown'}")
-        print(f"Description: {company_info.company_description}")
-    else:
-        print("Could not retrieve company information.")
+    async def main():
+        companies = ["Apple Inc.", "Microsoft", "Google", "A non-existent company", 
+                     "NVIDIA", "Tesla", "Amazon", "Meta"]
+        concurrency_limit = 5
+        
+        tasks = [search_company_info(company_name=name, template_service=template_service, 
+                                     model_name="gpt-4.1-mini") for name in companies]
+        results = await asyncio.gather(*tasks)
+        
+        company_info_map = {info.name: info for info in results if info}
+
+        for company_name in companies:
+            company_info = company_info_map.get(company_name)
+            if company_info:
+                print("\nCompany Information:")
+                print(f"  Name: {company_info.name}")
+                print(f"  Industry: {company_info.industry}")
+                print(f"  Size: {company_info.size} employees")
+                print(f"  Revenue: {company_info.revenue}")
+                print(f"  Location: {company_info.location}")
+                print(f"  Website: {company_info.website}")
+                print(f"  Founded: {company_info.founded_year or 'Unknown'}")
+                print(f"  Description: {company_info.description}")
+            else:
+                print(f"\nCould not retrieve company information for {company_name}.")
+
+    asyncio.run(main())

src/core/agents/utils/nodes.py

@@ -1,15 +1,19 @@
 import asyncio
-from typing import Dict, Optional
+from typing import Dict, Optional, Literal, List
 from langsmith import traceable
+import logging
+
 from src.core.agents.utils.state import AgentState
 from src.core.domain.resume import Resume
 from src.core.domain.job_description import JobDescription
-from src.core.domain.company_search import CompanySearchResponse
+from src.core.domain.company_search import CompanyInfo
 from src.infrastructure.components import llm_extractor
 from src.core.ports.secondary.llm_extractor import LLMExtractor
 from src.core.ports.secondary.template_service import TemplateService
 from src.core.agents.search_agents import search_company_info
 
+logger = logging.getLogger("core.agents.nodes")
+
 @traceable(run_type="llm")
 async def parse_resume_node(state: AgentState, extractor: LLMExtractor):
     resume = await extractor.parse_document(
@@ -33,14 +37,20 @@ async def parse_job_description_node(state: AgentState, extractor: LLMExtractor)
 @traceable(run_type="tool")
 async def search_company_info_node(state: AgentState, 
                                    template_service: TemplateService, 
-                                   model_name="gpt-4o") -> Dict[str, Optional[CompanySearchResponse]]:
-    """
-    Search for information about a company and return structured data.
-    """
-    company_names = [experience.company for experience in state["resume"].experiences]
-
-    tasks = [search_company_info(company, template_service, model_name) for company in company_names]
-    results = await asyncio.gather(*tasks)
-    company_info_list = [result for result in results if result is not None]
-    print(f"Retrieved information for {len(company_info_list)} out of {len(company_names)} companies.")
-    return {"company_info": company_info_list}
+                                   branch: Literal["resume", "job_desription"] = "job_description",
+                                   model_name="gpt-4.1-mini") -> Dict[str, Optional[CompanyInfo]]:
+    if branch == "resume":
+        companies = state["resume"].company_names
+    elif branch == "job_description":
+        companies = [state["job_description"].company_name]
+    else:
+        raise ValueError("This branch is not supported.")
+    
+    tasks = [search_company_info(company_name=company_name, 
+                                 model_name=model_name, 
+                                 template_service=template_service) for company_name in companies]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = {result.name: result for result in results if isinstance(result, CompanyInfo)}
+    logger.info(f"Retrieved information for branch {branch} with length {len(results)} companies.")
+
+    return {f"{branch}_company_info": results}

src/core/agents/utils/state.py

@@ -1,6 +1,6 @@
 from typing import List, Optional, TypedDict, Dict
 
-from src.core.domain.company_search import CompanySearchResponse
+from src.core.domain.company_search import CompanyInfo
 from src.core.domain.resume import Resume
 from src.core.domain.job_description import JobDescription
 from src.core.domain.resume_match import (
@@ -25,7 +25,8 @@ class AgentState(TypedDict):
     job_description: JobDescription
 
     # Enriched data
-    company_info: Optional[Dict[str, CompanySearchResponse]]
+    resume_company_info: Optional[Dict[str, CompanyInfo]]
+    job_description_company_info: Optional[Dict[str, CompanyInfo]]
 
     # Analysis results (populated by agents)
     skill_matches: Optional[List[SkillMatch]]

src/core/domain/company_search.py

@@ -1,17 +1,17 @@
 from pydantic import BaseModel, Field
 from typing import Optional
 
-class CompanySearchResponse(BaseModel):
-    company_name: str = Field(description="The name of the company")
-    company_description: Optional[str] = Field(None, description="A description of the company") 
-    company_website: Optional[str] = Field(None, description="The website of the company")
-    company_location: Optional[str] = Field(None, description="The location of the company")
-    company_industry: Optional[str] = Field(None, description="The industry of the company")
-    company_size: Optional[int] = Field(None, description="The approximate number of employees at the company")
-    company_revenue: Optional[str] = Field(None, description="The annual revenue of the company")
+class CompanyInfo(BaseModel):
+    name: str = Field(description="The name of the company")
+    description: Optional[str] = Field(None, description="A description of the company") 
+    website: Optional[str] = Field(None, description="The website of the company")
+    location: Optional[str] = Field(None, description="The location of the company")
+    industry: Optional[str] = Field(None, description="The industry of the company")
+    size: Optional[int] = Field(None, description="The approximate number of employees at the company")
+    revenue: Optional[str] = Field(None, description="The annual revenue of the company")
     founded_year: Optional[int] = Field(None, description="The year the company was founded")
 
 
 if __name__ == "__main__":
-    data_model = CompanySearchResponse.model_json_schema()
+    data_model = CompanyInfo.model_json_schema()
     print(data_model)

src/core/domain/resume.py

@@ -99,3 +99,7 @@ def parse_raw_json(cls, json_str: str) -> "Resume":
             return cls.model_validate_json(json_str)
         except json.JSONDecodeError as e:
             raise ValueError(f"Invalid JSON string: {str(e)}")
+
+    @property
+    def company_names(self) -> List[str]:
+        return [experience.company for experience in self.experiences]